<commit_before>extern crate hal;
extern crate arrayfire as af;
extern crate itertools;
extern crate rand;
#[macro_use] extern crate timeit;

use std::env;
use af::{Array, Dim4, Backend, DType};
use itertools::Zip;
use rand::distributions::{IndependentSample, Range};

use hal::{utils, activations, initializations, loss};
use hal::layer;
use hal::layer::{Layer};
use hal::params::{DenseGenerator, RNNGenerator, ParamManager};
use hal::device::{DeviceManagerFactory, Device};
use hal::error::HALError;

//todo: move all these tests into separate modules

/// test derivatives
fn verify_derivative<F>(ufunc: F, name: &str)
  where F : Fn(&Array) -> Array
{
  println!("\ngradient testing {}...", name);
  let dims = Dim4::new(&[1, 1, 1, 1]);

  // generate a random number between (-1, 1)
  let between = Range::new(-1f32, 1f32); // utils::constant needs an f32
  let mut rng = rand::thread_rng();
  let rnd_num = between.ind_sample(&mut rng);

  // build a constant r^1 array
  let x = utils::constant(dims, DType::F64, rnd_num);

  // get the original activation and the symbolic gradient
  let activated = ufunc(&x);
  let grad = activations::get_derivative(name, &activated).unwrap();

  // run the algorithm on non-smooth function based this is a single func
  // [ie non-chained] and thus should be almost exact
  utils::verify_gradient_kinks(|i| {
    let rhs = ufunc(&i);
    let v = utils::array_to_vec(&rhs);
    v[0]
  }, &x, 1e-5, &grad).unwrap();
}

#[test]
fn lrelu_gradient() {
  verify_derivative(activations::lrelu
                    , "lrelu");
}

#[test]
fn tanh_gradient() {
  verify_derivative(activations::tanh
                    , "tanh");
}

#[test]
fn sigmoid_gradient() {
  verify_derivative(activations::sigmoid
                    , "sigmoid");
}

// todo: why is this broken? do you need softmax + xentropy?
// #[test]
// fn softmax_gradient() {
//   verify_derivative(activations::softmax
//                     , "softmax"
//                     , false);
//}

#[test]
fn relu_gradient() {
  verify_derivative(activations::relu
                    , "relu");
}

#[test]
fn ones_gradient() {
  verify_derivative(activations::ones
                    , "ones");
}

/// test unitary functions
fn verify_func<F>(ufunc: F, name: &str, truth: &[f32])
  where F : Fn(&Array) -> Array
{
  env::set_var("rust_test_threads", "1");
  println!("\ntesting unitary function {}...", name);
  let dims = Dim4::new(&[5, 1, 1, 1]);
  let x = Array::new::<f32>(&[-1.0, 0.0, 1.0, 2.0, 3.0], dims);

  // verify with l2 loss
  let x_t = Array::new::<f32>(truth, dims);
  let l2 = loss::get_loss("l2", &ufunc(&x), &x_t).unwrap();
  assert!(l2 <= 1e-4, "l2 loss of {} is higher than expected: {}", name, l2);
}

#[test]
fn tanh(){
  verify_func(activations::tanh
              , "tanh"
              , &[-0.7616, 0.0000, 0.7616, 0.9640, 0.9951]);
}

#[test]
fn sigmoid(){
  verify_func(activations::sigmoid
              , "sigmoid"
              , &[0.2689, 0.5000, 0.7311, 0.8808, 0.9526]);
}

#[test]
fn relu(){
  verify_func(activations::relu
              , "relu"
              , &[0.0, 0.0, 1.0, 2.0, 3.0]);
}

#[test]
fn lrelu(){
  verify_func(activations::lrelu
              , "lrelu"
              , &[-0.01, 0.0, 1.0, 2.0, 3.0]);
}

#[test]
fn softmax(){
  verify_func(activations::softmax
              , "softmax"
              , &[0.0117, 0.0317, 0.0861, 0.2341, 0.6364]);
}

#[test]
fn ones(){
  verify_func(activations::ones
              , "ones"
              , &[-1.0, 0.0, 1.0, 2.0, 3.0]);
}

/// test losses
#[test]
fn cross_entropy_softmax(){
  println!("\ntesting cross-entropy soft max...");
  let dims = Dim4::new(&[5, 1, 1, 1]);
  let x = Array::new(&[-0.01, 0.00, 1.10, 2.20, 3.15], dims);
  let target = Array::new(&[1.0, 0.00, 0.00, 0.00, 0.00], dims);
  let x_t = 4.9980f32;

  let x_pred = loss::get_loss("cross_entropy"
                              , &activations::softmax(&x)
                              , &target).unwrap();
  let l2 = x_t * x_t - x_pred * x_pred;
  assert!(l2 < 1e-7
          , "cross-entropy-softmax loss is more than expected {} vs {} => {}"
          , x_t, x_pred, l2);
}

/// helper to build a layer
pub fn layer_builder<F>(layer_type: &str, idims: Dim4, odims: Dim4, loss: &str
                        , eps: f64, activation: &str, w_init: &str, b_init: &str, mut f: F)
  where F: FnMut(&ParamManager, Box<Layer>)
{
  // [batch_size, input_size, temporal_size, 1]
  let input_size: usize = idims[1] as usize;
  let output_size: usize = odims[1] as usize;
  let temporal_size: usize = idims[2] as usize;

  // add a param manager, a device manager, a device
  let mut param_manager = ParamManager::default();
  let device_manager = DeviceManagerFactory::new();
  let device = Device{backend: Backend::DEFAULT, id: 0};

  // add the layer type
  let layer: Box<Layer> = match layer_type {
    "Dense" => Box::new(layer::Dense {
      input_size: input_size,
      output_size: output_size,
    }),
    "rnn"  => Box::new(layer::RNN {
      input_size: input_size,
      output_size: output_size,
    }),
    //todo: lstm, etc
    _      => panic!("unknown layer type specified"),
  };

  // push it into the param manager
  match layer_type {
    "Dense" => param_manager.add_dense::<f64>(device_manager, device
                                              , input_size, output_size
                                              , activation
                                              , w_init
                                              , b_init),
    "rnn"  => param_manager.add_rnn::<f64>(device_manager, device
                                             , input_size, output_size
                                             , activation
                                             , w_init, w_init // just dupe it
                                             , b_init),
  //todo: lstm, etc
    _      => panic!("unknown layer type specified"),
  };

  // run the closure
  f(&mut param_manager, layer);
}

/// test forward pass for layers
pub fn layer_forward_helper(layer_type: &str, idims: Dim4, odims: Dim4, loss: &str
                            , eps: f64, activation: &str, w_init: &str, b_init: &str
                            , inputs_vec: Vec<f64>, targets_vec: Vec<f64>)
{
  //env::set_var("af_disable_graphics", "1"); // glfw crashes otherwise
  println!("testing {} layer with {} acivation for forward pass..."
           , layer_type, activation);
  let x = Array::new::<f64>(&inputs_vec[..], idims);
  let targets = Array::new::<f64>(&targets_vec[..], odims);

  layer_builder(layer_type, idims, odims, loss
                , eps, activation, w_init, b_init, |param_manager, layer|
                {
                  // run a forward pass and verify it is within tolerance
                  let params = param_manager.get_params(0);
                  let activ = layer.forward(params.clone(), &x.clone());

                  let loss_activ = loss::get_loss(loss, &activ, &targets).unwrap();
                  assert!(loss_activ < 1e-9
                          , "forward pass verification failed, error = {}"
                          , loss_activ);
                });
}

/// test layers gradients
pub fn layer_backward_helper(layer_type: &str, idims: Dim4, odims: Dim4, loss: &str
                             , eps: f64, activation: &str, w_init: &str, b_init: &str)
{
  //env::set_var("af_disable_graphics", "1"); // glfw crashes otherwise
  // [batch_size, input_size, temporal_size, 1]
  let input_size: usize = idims[1] as usize;
  let output_size: usize = odims[1] as usize;
  let temporal_size: usize = idims[2] as usize;

  let x = initializations::uniform::<f64>(idims, 0.5f32);
  let targets = match activation {
    "softmax" => {
      // randomly pick one of K indexes to set to 1
      let mut v: Vec<f64> = vec![0f64; output_size];
      let between = Range::new(0usize, output_size as usize);
      let mut rng = rand::thread_rng();
      let rnd_index = between.ind_sample(&mut rng);
      v[rnd_index] = 1f64;

      // build an array
      utils::vec_to_array::<f64>(v, odims)
    },

    _ => initializations::uniform::<f64>(idims, 0.5f32),
  };

  layer_builder(layer_type, idims, odims, loss
                , eps, activation, w_init, b_init, |param_manager, layer|
                {
                  // run a forward and then bkwd pass to extract the gradients
                  let params = param_manager.get_params(0);
                  let activ = layer.forward(params.clone(), &x.clone());
                  let delta = loss::get_loss_derivative(loss, &activ, &targets).unwrap();
                  layer.backward(params.clone(), &delta);
                  let grads = param_manager.get_all_deltas();
                  let num_params = param_manager.num_arrays(0);
                  let bkp_input = param_manager.get_inputs(0);
                  let bkp_output = param_manager.get_inputs(0);
                  let bkp_recur = param_manager.get_recurrences(0);

                  // iterate over all arrays and grads and run gradient checking
                  for (arr, grad, ind) in Zip::new((param_manager.get_all_arrays().iter() // weights + biases
                                                    , grads                               // tabulated gradients
                                                    , 0..num_params))                     // param index iterator
                  {
                    let arr_bkp: Array = arr.copy(); // keep a backup

                    // do the gradient check specific to the activation type
                    match activations::is_smooth(activation) {
                      false => utils::verify_gradient_kinks(|i| {
                        // run forward pass using the modified array
                        let p = params.clone();
                        p.lock().unwrap().current_unroll = 0;
                        param_manager.set_array_from_index(i.clone(), ind);
                        param_manager.set_inputs(0, bkp_input.clone());
                        param_manager.set_outputs(0, bkp_output.clone());
                        param_manager.set_recurrences(0, bkp_recur.clone());
                        let fwd_pass = layer.forward(params.clone(), &x.clone());
                        loss::get_loss(loss, &fwd_pass, &targets).unwrap() as f64
                      }, &arr_bkp, eps, &grad).unwrap(),
                      true  => utils::verify_gradient_smooth(|i| {
                        // run forward pass using the modified array
                        let p = params.clone();
                        p.lock().unwrap().current_unroll = 0;
                        param_manager.set_array_from_index(i.clone(), ind);
                        param_manager.set_inputs(0, bkp_input.clone());
                        param_manager.set_outputs(0, bkp_output.clone());
                        param_manager.set_recurrences(0, bkp_recur.clone());
                        let fwd_pass = layer.forward(params.clone(), &x.clone());
                        loss::get_loss(loss, &fwd_pass, &targets).unwrap() as f64
                      }, &arr_bkp, eps, &grad).unwrap(),
                    };
                  }
                });
}

#[test]
fn dense_forward(){
  timeit!({
    let idims = Dim4::new(&[1, 5, 1, 1]);
    let odims = Dim4::new(&[1, 5, 1, 1]);
    layer_forward_helper("Dense", idims, odims, "l2", 1e-4
                         , "linear"                                      // activation
                         , "ones"                                        // weight init
                         , "zeros"                                       // bias init
                         , vec![-0.01, 0.00, 1.10, 2.20, 3.15]           //input
                         , vec![6.4400, 6.4400,6.4400, 6.4400, 6.4400]); //target
  });
}

#[test]
fn dense_backward() {
  timeit! ({
    let idims = Dim4::new(&[1, 5, 1, 1]);
    let odims = Dim4::new(&[1, 5, 1, 1]);
    layer_backward_helper("Dense", idims, odims
                          , "l2"              // loss
                          , 1e-4              // eps for numerical grad
                          , "tanh"            // activation
                          , "glorot_uniform"  // weight init
                          , "zeros");         // bias init
  });
}

#[test]
fn rnn_forward(){
  timeit!({
    let idims = Dim4::new(&[1, 5, 1, 1]);
    let odims = Dim4::new(&[1, 5, 1, 1]);
    layer_forward_helper("rnn", idims, odims, "l2", 1e-4
                         , "linear"                                      // activation
                         , "ones"                                        // weight init
                         , "zeros"                                       // bias init
                         , vec![-0.01, 0.00, 1.10, 2.20, 3.15]           //input
                         , vec![6.4400, 6.4400,6.4400, 6.4400, 6.4400]); //target
  });
}

#[test]
fn rnn_backward() {
  timeit! ({
    let idims = Dim4::new(&[1, 5, 1, 1]); // single time slice
    let odims = Dim4::new(&[1, 5, 1, 1]); // single time slice
    layer_backward_helper("rnn", idims, odims
                          , "l2"              // loss
                          , 1e-4              // eps for numerical grad
                          , "tanh"            // activation
                          , "glorot_uniform"  // weight init
                          , "zeros");         // bias init
  });
}
<commit_msg>update test to work with a sample hidden layer<commit_after>extern crate hal;
extern crate arrayfire as af;
extern crate itertools;
extern crate rand;
#[macro_use] extern crate timeit;

use std::env;
use af::{Array, Dim4, Backend, DType};
use itertools::Zip;
use rand::distributions::{IndependentSample, Range};

use hal::{utils, activations, initializations, loss};
use hal::layer;
use hal::layer::{Layer};
use hal::params::{DenseGenerator, RNNGenerator, ParamManager};
use hal::device::{DeviceManagerFactory, Device};
use hal::error::HALError;

//todo: move all these tests into separate modules

/// test derivatives
fn verify_derivative<F>(ufunc: F, name: &str)
  where F : Fn(&Array) -> Array
{
  println!("\ngradient testing {}...", name);
  let dims = Dim4::new(&[1, 1, 1, 1]);

  // generate a random number between (-1, 1)
  let between = Range::new(-1f32, 1f32); // utils::constant needs an f32
  let mut rng = rand::thread_rng();
  let rnd_num = between.ind_sample(&mut rng);

  // build a constant r^1 array
  let x = utils::constant(dims, DType::F64, rnd_num);

  // get the original activation and the symbolic gradient
  let activated = ufunc(&x);
  let grad = activations::get_derivative(name, &activated).unwrap();

  // run the algorithm on non-smooth function based this is a single func
  // [ie non-chained] and thus should be almost exact
  utils::verify_gradient_kinks(|i| {
    let rhs = ufunc(&i);
    let v = utils::array_to_vec(&rhs);
    v[0]
  }, &x, 1e-5, &grad).unwrap();
}

#[test]
fn lrelu_gradient() {
  verify_derivative(activations::lrelu
                    , "lrelu");
}

#[test]
fn tanh_gradient() {
  verify_derivative(activations::tanh
                    , "tanh");
}

#[test]
fn sigmoid_gradient() {
  verify_derivative(activations::sigmoid
                    , "sigmoid");
}

// todo: why is this broken? do you need softmax + xentropy?
// #[test]
// fn softmax_gradient() {
//   verify_derivative(activations::softmax
//                     , "softmax"
//                     , false);
//}

#[test]
fn relu_gradient() {
  verify_derivative(activations::relu
                    , "relu");
}

#[test]
fn ones_gradient() {
  verify_derivative(activations::ones
                    , "ones");
}

/// test unitary functions
fn verify_func<F>(ufunc: F, name: &str, truth: &[f32])
  where F : Fn(&Array) -> Array
{
  env::set_var("rust_test_threads", "1");
  println!("\ntesting unitary function {}...", name);
  let dims = Dim4::new(&[5, 1, 1, 1]);
  let x = Array::new::<f32>(&[-1.0, 0.0, 1.0, 2.0, 3.0], dims);

  // verify with l2 loss
  let x_t = Array::new::<f32>(truth, dims);
  let l2 = loss::get_loss("l2", &ufunc(&x), &x_t).unwrap();
  assert!(l2 <= 1e-4, "l2 loss of {} is higher than expected: {}", name, l2);
}

#[test]
fn tanh(){
  verify_func(activations::tanh
              , "tanh"
              , &[-0.7616, 0.0000, 0.7616, 0.9640, 0.9951]);
}

#[test]
fn sigmoid(){
  verify_func(activations::sigmoid
              , "sigmoid"
              , &[0.2689, 0.5000, 0.7311, 0.8808, 0.9526]);
}

#[test]
fn relu(){
  verify_func(activations::relu
              , "relu"
              , &[0.0, 0.0, 1.0, 2.0, 3.0]);
}

#[test]
fn lrelu(){
  verify_func(activations::lrelu
              , "lrelu"
              , &[-0.01, 0.0, 1.0, 2.0, 3.0]);
}

#[test]
fn softmax(){
  verify_func(activations::softmax
              , "softmax"
              , &[0.0117, 0.0317, 0.0861, 0.2341, 0.6364]);
}

#[test]
fn ones(){
  verify_func(activations::ones
              , "ones"
              , &[-1.0, 0.0, 1.0, 2.0, 3.0]);
}

/// test losses
#[test]
fn cross_entropy_softmax(){
  println!("\ntesting cross-entropy soft max...");
  let dims = Dim4::new(&[5, 1, 1, 1]);
  let x = Array::new(&[-0.01, 0.00, 1.10, 2.20, 3.15], dims);
  let target = Array::new(&[1.0, 0.00, 0.00, 0.00, 0.00], dims);
  let x_t = 4.9980f32;

  let x_pred = loss::get_loss("cross_entropy"
                              , &activations::softmax(&x)
                              , &target).unwrap();
  let l2 = x_t * x_t - x_pred * x_pred;
  assert!(l2 < 1e-7
          , "cross-entropy-softmax loss is more than expected {} vs {} => {}"
          , x_t, x_pred, l2);
}

/// helper to build a layer
pub fn layer_builder<F>(layer_type: &str, idims: Dim4, odims: Dim4, loss: &str
                        , eps: f64, activation: &str, w_init: &str, b_init: &str, mut f: F)
  where F: FnMut(&ParamManager, Box<Layer>)
{
  // [batch_size, input_size, temporal_size, 1]
  let batch_size: usize = idims[0] as usize;
  let input_size: usize = idims[1] as usize;
  let output_size: usize = odims[1] as usize;
  let temporal_size: usize = idims[2] as usize;

  // add a param manager, a device manager, a device
  let mut param_manager = ParamManager::default();
  let device_manager = DeviceManagerFactory::new();
  let device = Device{backend: Backend::DEFAULT, id: 0};

  // add the layer type
  let layer: Box<Layer> = match layer_type {
    "Dense" => Box::new(layer::Dense {
      input_size: input_size,
      output_size: output_size,
    }),
    "rnn"  => Box::new(layer::RNN {
      input_size: input_size,
      output_size: output_size,
    }),
    //todo: lstm, etc
    _      => panic!("unknown layer type specified"),
  };

  // push it into the param manager
  match layer_type {
    "Dense" => param_manager.add_dense::<f64>(device_manager, device
                                              , input_size, output_size
                                              , activation
                                              , w_init
                                              , b_init),
    "rnn"  => {
      param_manager.add_rnn::<f64>(device_manager, device
                                   , input_size, output_size
                                   , activation
                                   , w_init, w_init // just dupe it
                                   , b_init);
      // make it such that we are within an unrolling
      let hdims = Dim4::new(&[batch_size as u64, output_size as u64, 1, 1]);
      let h_t = utils::constant(hdims, DType::F64, 0.5f32);
      param_manager.set_recurrences(0, vec![h_t]);
    }
  //todo: lstm, etc
    _      => panic!("unknown layer type specified"),
  };

  // run the closure
  f(&mut param_manager, layer);
}

/// test forward pass for layers
pub fn layer_forward_helper(layer_type: &str, idims: Dim4, odims: Dim4, loss: &str
                            , eps: f64, activation: &str, w_init: &str, b_init: &str
                            , inputs_vec: Vec<f64>, targets_vec: Vec<f64>)
{
  //env::set_var("af_disable_graphics", "1"); // glfw crashes otherwise
  println!("testing {} layer with {} acivation for forward pass..."
           , layer_type, activation);
  let x = Array::new::<f64>(&inputs_vec[..], idims);
  let targets = Array::new::<f64>(&targets_vec[..], odims);

  layer_builder(layer_type, idims, odims, loss
                , eps, activation, w_init, b_init, |param_manager, layer|
                {
                  // run a forward pass and verify it is within tolerance
                  let params = param_manager.get_params(0);
                  let activ = layer.forward(params.clone(), &x.clone());

                  let loss_activ = loss::get_loss(loss, &activ, &targets).unwrap();
                  assert!(loss_activ < 1e-9
                          , "forward pass verification failed, error = {}"
                          , loss_activ);
                });
}

/// test layers gradients
pub fn layer_backward_helper(layer_type: &str, idims: Dim4, odims: Dim4, loss: &str
                             , eps: f64, activation: &str, w_init: &str, b_init: &str)
{
  //env::set_var("af_disable_graphics", "1"); // glfw crashes otherwise
  // [batch_size, input_size, temporal_size, 1]
  let input_size: usize = idims[1] as usize;
  let output_size: usize = odims[1] as usize;
  let temporal_size: usize = idims[2] as usize;

  let x = initializations::uniform::<f64>(idims, 0.5f32);
  let targets = match activation {
    "softmax" => {
      // randomly pick one of K indexes to set to 1
      let mut v: Vec<f64> = vec![0f64; output_size];
      let between = Range::new(0usize, output_size as usize);
      let mut rng = rand::thread_rng();
      let rnd_index = between.ind_sample(&mut rng);
      v[rnd_index] = 1f64;

      // build an array
      utils::vec_to_array::<f64>(v, odims)
    },

    _ => initializations::uniform::<f64>(idims, 0.5f32),
  };

  layer_builder(layer_type, idims, odims, loss
                , eps, activation, w_init, b_init, |param_manager, layer|
                {
                  // run a forward and then bkwd pass to extract the gradients
                  let params = param_manager.get_params(0);
                  let activ = layer.forward(params.clone(), &x.clone());
                  let delta = loss::get_loss_derivative(loss, &activ, &targets).unwrap();
                  layer.backward(params.clone(), &delta);
                  let grads = param_manager.get_all_deltas();
                  let num_params = param_manager.num_arrays(0);
                  let bkp_input = param_manager.get_inputs(0);
                  let bkp_output = param_manager.get_inputs(0);
                  let bkp_recur = param_manager.get_recurrences(0);

                  // iterate over all arrays and grads and run gradient checking
                  for (arr, grad, ind) in Zip::new((param_manager.get_all_arrays().iter() // weights + biases
                                                    , grads                               // tabulated gradients
                                                    , 0..num_params))                     // param index iterator
                  {
                    let arr_bkp: Array = arr.copy(); // keep a backup

                    // do the gradient check specific to the activation type
                    match activations::is_smooth(activation) {
                      false => utils::verify_gradient_kinks(|i| {
                        // run forward pass using the modified array
                        let p = params.clone();
                        p.lock().unwrap().current_unroll = 0;
                        param_manager.set_array_from_index(i.clone(), ind);
                        param_manager.set_inputs(0, bkp_input.clone());
                        param_manager.set_outputs(0, bkp_output.clone());
                        param_manager.set_recurrences(0, bkp_recur.clone());
                        let fwd_pass = layer.forward(params.clone(), &x.clone());
                        loss::get_loss(loss, &fwd_pass, &targets).unwrap() as f64
                      }, &arr_bkp, eps, &grad).unwrap(),
                      true  => utils::verify_gradient_smooth(|i| {
                        // run forward pass using the modified array
                        let p = params.clone();
                        p.lock().unwrap().current_unroll = 0;
                        param_manager.set_array_from_index(i.clone(), ind);
                        param_manager.set_inputs(0, bkp_input.clone());
                        param_manager.set_outputs(0, bkp_output.clone());
                        param_manager.set_recurrences(0, bkp_recur.clone());
                        let fwd_pass = layer.forward(params.clone(), &x.clone());
                        loss::get_loss(loss, &fwd_pass, &targets).unwrap() as f64
                      }, &arr_bkp, eps, &grad).unwrap(),
                    };
                  }
                });
}

#[test]
fn dense_forward(){
  timeit!({
    let idims = Dim4::new(&[1, 5, 1, 1]);
    let odims = Dim4::new(&[1, 5, 1, 1]);
    layer_forward_helper("Dense", idims, odims, "l2", 1e-4
                         , "linear"                                      // activation
                         , "ones"                                        // weight init
                         , "zeros"                                       // bias init
                         , vec![-0.01, 0.00, 1.10, 2.20, 3.15]           //input
                         , vec![6.4400, 6.4400,6.4400, 6.4400, 6.4400]); //target
  });
}

#[test]
fn dense_backward() {
  timeit! ({
    let idims = Dim4::new(&[1, 5, 1, 1]);
    let odims = Dim4::new(&[1, 5, 1, 1]);
    layer_backward_helper("Dense", idims, odims
                          , "l2"              // loss
                          , 1e-4              // eps for numerical grad
                          , "tanh"            // activation
                          , "glorot_uniform"  // weight init
                          , "zeros");         // bias init
  });
}

#[test]
fn rnn_forward(){
  timeit!({
    let idims = Dim4::new(&[1, 5, 1, 1]);
    let odims = Dim4::new(&[1, 5, 1, 1]);
    layer_forward_helper("rnn", idims, odims, "l2", 1e-4
                         , "linear"                                      // activation
                         , "ones"                                        // weight init
                         , "zeros"                                       // bias init
                         , vec![-0.01, 0.00, 1.10, 2.20, 3.15]           //input
                         , vec![ 8.94000053, 8.94000053, 8.94000053, 8.94000053, 8.94000053]); //target
  });
}

#[test]
fn rnn_backward() {
  timeit! ({
    let idims = Dim4::new(&[1, 5, 1, 1]); // single time slice
    let odims = Dim4::new(&[1, 5, 1, 1]); // single time slice
    layer_backward_helper("rnn", idims, odims
                          , "l2"              // loss
                          , 1e-4              // eps for numerical grad
                          , "tanh"            // activation
                          , "glorot_uniform"  // weight init
                          , "zeros");         // bias init
  });
}
<|endoftext|>


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

//! # Command Parser
//!
//! The `commands::parser` module provides:
//!
//! * Functionality for building a tree of commands based on their names.
//! * Parsing tokenized input to select a command. Input tokenization can
//!   be performed by the `commands::tokenizer` module.
//! * Validating parameters.
//! * Performing completion on commands and parameters.
//!
//! The parser is still a work in progress but is rapidly advancing.

pub mod nodes;
pub mod completion;
pub mod builder;

use std::collections::HashMap;
use std::error::Error;
use std::fmt;
use std::rc::Rc;
use parser::nodes::*;
use tokenizer::{Token, TokenType};
use parser::completion::{Complete, Completion};

/// Command parser
///
/// The lifetime parameter `'p` refers to the lifetime of the
/// parser itself, the root node of the command tree that was
/// used to create this parser, as well as the tokens passed
/// into the parser.
///
/// ```
/// use commands::parser::nodes::RootNode;
/// use commands::parser::Parser;
///
/// let root = RootNode::new();
/// let mut parser = Parser::new(root);
/// ```
///
/// The parser is constructed as a `mut`able object as most of
/// the methods on it will modify its state.
pub struct Parser<'p> {
    current_node: Rc<Node>,
    /// The nodes which have been accepted during `parse` or `advance`.
    pub nodes: Vec<Rc<Node>>,
    /// The tokens which have been accepted during `parse` or `advance`.
    pub tokens: Vec<Token<'p>>,
    commands: Vec<Rc<CommandNode>>,
    parameters: HashMap<String, String>,
}

impl<'p> Parser<'p> {
    /// Construct a parser with a root node.
    pub fn new(initial_node: Rc<RootNode>) -> Parser<'p> {
        Parser {
            current_node: initial_node,
            nodes: vec![],
            tokens: vec![],
            commands: vec![],
            parameters: HashMap::new(),
        }
    }

    /// Given an optional token, get the possible valid completions
    /// for the current parser state.
    ///
    /// Possible completions are successors of the current node which
    /// are not `hidden`, are `acceptable`, and which match the token,
    /// if one has been provided.
    ///
    /// Nodes may customize the `Complete` trait to customize the
    /// `Completion` and `CompletionOption`s which are generated
    /// for that node.
    ///
    /// Each valid successor node will have one `Completion` in the
    /// result vector. Each `Completion` will have one or more
    /// `CompletionOption` for each valid way that the value may be
    /// entered.
    pub fn complete(&self, token: Option<Token<'p>>) -> Vec<Completion> {
        self.current_node
            .successors()
            .into_iter()
            .filter(|n| {
                // To be a possible completion, the node should not be
                // hidden, it should be acceptable, and if there's a token,
                // it should be a valid match for the node.
                !n.hidden() && n.acceptable(self) &&
                if let Some(t) = token {
                    n.matches(self, t)
                } else {
                    true
                }
            })
            .map(|n| n.complete(token))
            .collect::<Vec<_>>()
    }

    /// Parse a vector of tokens, advancing through the
    /// node hierarchy.
    ///
    /// ```
    /// use commands::parser::nodes::RootNode;
    /// use commands::parser::Parser;
    /// use commands::tokenizer::tokenize;
    ///
    /// let root = RootNode::new();
    /// let mut parser = Parser::new(root);
    ///
    /// if let Ok(tokens) = tokenize("show interface") {
    ///     parser.parse(tokens);
    /// }
    /// ```
    pub fn parse(&mut self, tokens: Vec<Token<'p>>) -> Result<(), ParseError<'p>> {
        for token in tokens {
            match token.token_type {
                TokenType::Invalid => unreachable!(),
                TokenType::Whitespace => {}
                TokenType::Word => try!(self.advance(token)),
            }
        }
        Ok(())
    }

    /// Parse a single token, advancing through the node hierarchy.
    pub fn advance(&mut self, token: Token<'p>) -> Result<(), ParseError<'p>> {
        // We clone the current node so that it doesn't stay borrowed
        // and break things when we try to modify it below.
        let cn = self.current_node.clone();
        let matches = cn.successors()
                        .into_iter()
                        .filter(|n| n.acceptable(self) && n.matches(self, token))
                        .collect::<Vec<_>>();
        match matches.len() {
            1 => {
                let matching_node = matches[0];
                matching_node.accept(self, token);
                self.current_node = matching_node.clone();
                self.nodes.push(matching_node.clone());
                self.tokens.push(token);
                Ok(())
            }
            0 => Err(ParseError::NoMatches(token)),
            _ => Err(ParseError::AmbiguousMatch(token)),
        }
    }

    /// Execute the command that has been accepted by the parser.
    ///
    /// * XXX: This should be returning a Result probably.
    pub fn execute(&self) {
        if !self.commands.is_empty() {
            unimplemented!();
            // self.commands[0].execute(self)
        }
    }

    /// Verify that the parser is in a valid state with
    /// respect to having accepted a command and all
    /// required parameters.
    pub fn verify(&self) -> Result<(), VerifyError> {
        if self.commands.is_empty() {
            return Err(VerifyError::NoCommandAccepted);
        } else {
            for expected in self.commands[0].parameters() {
                if expected.required() && !self.parameters.contains_key(expected.name()) {
                    return Err(VerifyError::MissingParameter(expected.name().clone()));
                }
            }
        }
        Ok(())
    }
}

/// Errors that calling `parse` on the `Parser` can raise.
#[derive(Clone,Debug)]
pub enum ParseError<'t> {
    /// The parser is in an invalid state.
    InvalidState,
    /// There were no matches for the token.
    NoMatches(Token<'t>),
    /// There was more than 1 possible match for the token.
    AmbiguousMatch(Token<'t>), // XXX: One day, add: Vec<&'p Rc<Node>>),
}

impl<'t> Error for ParseError<'t> {
    fn description(&self) -> &str {
        match *self {
            ParseError::InvalidState => "Invalid state.",
            ParseError::NoMatches(_) => "No match.",
            ParseError::AmbiguousMatch(_) => "Ambiguous match.",
        }
    }
}

impl<'t> fmt::Display for ParseError<'t> {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        self.description().fmt(f)
    }
}

/// Errors that calling `verify` on the `Parser` can raise.
#[derive(Clone,Debug)]
pub enum VerifyError {
    /// No command has been accepted by the parser.
    NoCommandAccepted,
    /// A required parameter is missing.
    MissingParameter(String),
}

impl Error for VerifyError {
    fn description(&self) -> &str {
        match *self {
            VerifyError::NoCommandAccepted => "No command has been accepted by the parser.",
            VerifyError::MissingParameter(_) => "A required parameter is missing.",
        }
    }
}

impl fmt::Display for VerifyError {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        self.description().fmt(f)
    }
}

/// Can this node be accepted in the current parser state?
pub trait Acceptable {
    /// Can this node be accepted in the current parser state?
    fn acceptable(&self, parser: &Parser) -> bool;
}

impl Acceptable for Rc<Node> {
    /// By default, a node can be accepted when it hasn't been seen yet.
    fn acceptable(&self, parser: &Parser) -> bool {
        !parser.nodes.contains(self)
    }
}

impl Acceptable for RepeatableNode {
    /// A repeatable node can be accepted.
    fn acceptable(&self, _parser: &Parser) -> bool {
        if self.repeatable() {
            return true;
        }
        unimplemented!()
        // This should check nodes.contains, but then go on to check
        // for a repeat marker and whether or not that's been seen.
    }
}

/// Define whether or not a node matches a token.
pub trait Matches {
    /// Does this node match the specified `token`?
    fn matches(&self, parser: &Parser, token: Token) -> bool;
}

impl Matches for Node {
    /// By default, a node matches a `token` when the name of the
    /// node starts with the `token`.
    fn matches(&self, _parser: &Parser, token: Token) -> bool {
        self.name().starts_with(token.text)
    }
}

/// Define the behavior of a node when it is accepted by the parser.
pub trait Accept {
    /// Accept this node with the given `token` as data.
    ///
    /// This is where parameters are stored, commands added.
    fn accept<'p>(&self, parser: &mut Parser<'p>, token: Token);
}

impl Accept for Node {
    /// By default, nothing needs to happen for `accept`.
    fn accept<'p>(&self, _parser: &mut Parser<'p>, _token: Token) {}
}

impl Accept for Rc<CommandNode> {
    /// Record this command.
    fn accept<'p>(&self, parser: &mut Parser<'p>, _token: Token) {
        if let Some(_) = self.handler() {
            parser.commands.push(self.clone())
        }
    }
}

impl Accept for ParameterNode {
    /// Record this parameter value.
    fn accept<'p>(&self, parser: &mut Parser<'p>, token: Token) {
        if self.repeatable() {
            unimplemented!();
        } else {
            parser.parameters.insert(self.name().clone(), token.text.to_string());
        }
    }
}

#[cfg(test)]
mod test {
    use super::nodes::*;
    use super::*;

    #[test]
    #[should_panic]
    fn verify_signals_no_command() {
        let root = RootNode::new();
        let parser = Parser::new(root);
        match parser.verify() {
            Err(VerifyError::NoCommandAccepted) => panic!(),
            _ => {}
        }
    }
}
<commit_msg>Update parser docs.<commit_after>// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! # Command Parser
//!
//! The `commands::parser` module provides:
//!
//! * Functionality for building a tree of commands based on their names.
//! * Parsing tokenized input to select a command. Input tokenization can
//!   be performed by the `commands::tokenizer` module.
//! * Validating parameters.
//! * Performing completion on commands and parameters.
//!
//! The command parser consists of two important things:
//!
//! * A tree that represents the available commands and their arguments.
//!   This tree consists of instances of `Nodes` from the
//!   `commands::parser::nodes` module. Construction of this tree
//!   is done with the help of the `commands::parser::builder` module.
//! * A `Parser` that handles input and matches it against the command
//!   tree. This parser is intended to be short-lived and to just live
//!   for the duration of parsing and evaluating a single command line
//!   input.

pub mod nodes;
pub mod completion;
pub mod builder;

use std::collections::HashMap;
use std::error::Error;
use std::fmt;
use std::rc::Rc;
use parser::nodes::*;
use tokenizer::{Token, TokenType};
use parser::completion::{Complete, Completion};

/// Command parser
///
/// The lifetime parameter `'p` refers to the lifetime of the
/// parser itself, the root node of the command tree that was
/// used to create this parser, as well as the tokens passed
/// into the parser.
///
/// ```
/// use commands::parser::nodes::RootNode;
/// use commands::parser::Parser;
///
/// let root = RootNode::new();
/// let mut parser = Parser::new(root);
/// ```
///
/// The parser is constructed as a `mut`able object as most of
/// the methods on it will modify its state.
pub struct Parser<'p> {
    current_node: Rc<Node>,
    /// The nodes which have been accepted during `parse` or `advance`.
    pub nodes: Vec<Rc<Node>>,
    /// The tokens which have been accepted during `parse` or `advance`.
    pub tokens: Vec<Token<'p>>,
    commands: Vec<Rc<CommandNode>>,
    parameters: HashMap<String, String>,
}

impl<'p> Parser<'p> {
    /// Construct a parser with a root node.
    pub fn new(initial_node: Rc<RootNode>) -> Parser<'p> {
        Parser {
            current_node: initial_node,
            nodes: vec![],
            tokens: vec![],
            commands: vec![],
            parameters: HashMap::new(),
        }
    }

    /// Given an optional token, get the possible valid completions
    /// for the current parser state.
    ///
    /// Possible completions are successors of the current node which
    /// are not `hidden`, are `acceptable`, and which match the token,
    /// if one has been provided.
    ///
    /// Nodes may customize the `Complete` trait to customize the
    /// `Completion` and `CompletionOption`s which are generated
    /// for that node.
    ///
    /// Each valid successor node will have one `Completion` in the
    /// result vector. Each `Completion` will have one or more
    /// `CompletionOption` for each valid way that the value may be
    /// entered.
    pub fn complete(&self, token: Option<Token<'p>>) -> Vec<Completion> {
        self.current_node
            .successors()
            .into_iter()
            .filter(|n| {
                // To be a possible completion, the node should not be
                // hidden, it should be acceptable, and if there's a token,
                // it should be a valid match for the node.
                !n.hidden() && n.acceptable(self) &&
                if let Some(t) = token {
                    n.matches(self, t)
                } else {
                    true
                }
            })
            .map(|n| n.complete(token))
            .collect::<Vec<_>>()
    }

    /// Parse a vector of tokens, advancing through the
    /// node hierarchy.
    ///
    /// ```
    /// use commands::parser::nodes::RootNode;
    /// use commands::parser::Parser;
    /// use commands::tokenizer::tokenize;
    ///
    /// let root = RootNode::new();
    /// let mut parser = Parser::new(root);
    ///
    /// if let Ok(tokens) = tokenize("show interface") {
    ///     parser.parse(tokens);
    /// }
    /// ```
    pub fn parse(&mut self, tokens: Vec<Token<'p>>) -> Result<(), ParseError<'p>> {
        for token in tokens {
            match token.token_type {
                TokenType::Invalid => unreachable!(),
                TokenType::Whitespace => {}
                TokenType::Word => try!(self.advance(token)),
            }
        }
        Ok(())
    }

    /// Parse a single token, advancing through the node hierarchy.
    pub fn advance(&mut self, token: Token<'p>) -> Result<(), ParseError<'p>> {
        // We clone the current node so that it doesn't stay borrowed
        // and break things when we try to modify it below.
        let cn = self.current_node.clone();
        let matches = cn.successors()
                        .into_iter()
                        .filter(|n| n.acceptable(self) && n.matches(self, token))
                        .collect::<Vec<_>>();
        match matches.len() {
            1 => {
                let matching_node = matches[0];
                matching_node.accept(self, token);
                self.current_node = matching_node.clone();
                self.nodes.push(matching_node.clone());
                self.tokens.push(token);
                Ok(())
            }
            0 => Err(ParseError::NoMatches(token)),
            _ => Err(ParseError::AmbiguousMatch(token)),
        }
    }

    /// Execute the command that has been accepted by the parser.
    ///
    /// * XXX: This should be returning a Result probably.
    pub fn execute(&self) {
        if !self.commands.is_empty() {
            unimplemented!();
            // self.commands[0].execute(self)
        }
    }

    /// Verify that the parser is in a valid state with
    /// respect to having accepted a command and all
    /// required parameters.
    pub fn verify(&self) -> Result<(), VerifyError> {
        if self.commands.is_empty() {
            return Err(VerifyError::NoCommandAccepted);
        } else {
            for expected in self.commands[0].parameters() {
                if expected.required() && !self.parameters.contains_key(expected.name()) {
                    return Err(VerifyError::MissingParameter(expected.name().clone()));
                }
            }
        }
        Ok(())
    }
}

/// Errors that calling `parse` on the `Parser` can raise.
#[derive(Clone,Debug)]
pub enum ParseError<'t> {
    /// The parser is in an invalid state.
    InvalidState,
    /// There were no matches for the token.
    NoMatches(Token<'t>),
    /// There was more than 1 possible match for the token.
    AmbiguousMatch(Token<'t>), // XXX: One day, add: Vec<&'p Rc<Node>>),
}

impl<'t> Error for ParseError<'t> {
    fn description(&self) -> &str {
        match *self {
            ParseError::InvalidState => "Invalid state.",
            ParseError::NoMatches(_) => "No match.",
            ParseError::AmbiguousMatch(_) => "Ambiguous match.",
        }
    }
}

impl<'t> fmt::Display for ParseError<'t> {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        self.description().fmt(f)
    }
}

/// Errors that calling `verify` on the `Parser` can raise.
#[derive(Clone,Debug)]
pub enum VerifyError {
    /// No command has been accepted by the parser.
    NoCommandAccepted,
    /// A required parameter is missing.
    MissingParameter(String),
}

impl Error for VerifyError {
    fn description(&self) -> &str {
        match *self {
            VerifyError::NoCommandAccepted => "No command has been accepted by the parser.",
            VerifyError::MissingParameter(_) => "A required parameter is missing.",
        }
    }
}

impl fmt::Display for VerifyError {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        self.description().fmt(f)
    }
}

/// Can this node be accepted in the current parser state?
pub trait Acceptable {
    /// Can this node be accepted in the current parser state?
    fn acceptable(&self, parser: &Parser) -> bool;
}

impl Acceptable for Rc<Node> {
    /// By default, a node can be accepted when it hasn't been seen yet.
    fn acceptable(&self, parser: &Parser) -> bool {
        !parser.nodes.contains(self)
    }
}

impl Acceptable for RepeatableNode {
    /// A repeatable node can be accepted.
    fn acceptable(&self, _parser: &Parser) -> bool {
        if self.repeatable() {
            return true;
        }
        unimplemented!()
        // This should check nodes.contains, but then go on to check
        // for a repeat marker and whether or not that's been seen.
    }
}

/// Define whether or not a node matches a token.
pub trait Matches {
    /// Does this node match the specified `token`?
    fn matches(&self, parser: &Parser, token: Token) -> bool;
}

impl Matches for Node {
    /// By default, a node matches a `token` when the name of the
    /// node starts with the `token`.
    fn matches(&self, _parser: &Parser, token: Token) -> bool {
        self.name().starts_with(token.text)
    }
}

/// Define the behavior of a node when it is accepted by the parser.
pub trait Accept {
    /// Accept this node with the given `token` as data.
    ///
    /// This is where parameters are stored, commands added.
    fn accept<'p>(&self, parser: &mut Parser<'p>, token: Token);
}

impl Accept for Node {
    /// By default, nothing needs to happen for `accept`.
    fn accept<'p>(&self, _parser: &mut Parser<'p>, _token: Token) {}
}

impl Accept for Rc<CommandNode> {
    /// Record this command.
    fn accept<'p>(&self, parser: &mut Parser<'p>, _token: Token) {
        if let Some(_) = self.handler() {
            parser.commands.push(self.clone())
        }
    }
}

impl Accept for ParameterNode {
    /// Record this parameter value.
    fn accept<'p>(&self, parser: &mut Parser<'p>, token: Token) {
        if self.repeatable() {
            unimplemented!();
        } else {
            parser.parameters.insert(self.name().clone(), token.text.to_string());
        }
    }
}

#[cfg(test)]
mod test {
    use super::nodes::*;
    use super::*;

    #[test]
    #[should_panic]
    fn verify_signals_no_command() {
        let root = RootNode::new();
        let parser = Parser::new(root);
        match parser.verify() {
            Err(VerifyError::NoCommandAccepted) => panic!(),
            _ => {}
        }
    }
}
<|endoftext|>


<commit_before>use nom::{digit, hex_digit, oct_digit};
use std::{str, i32};

// Temporarly use a f64 to represent JS numbers
type JSNumber = f64;

// TODO: Implement proper parsing (and saving) of the number literal.

fn bytes_to_exponent(bytes: &[u8]) -> JSNumber {
    // TODO: Fix this unwrap
    let string = str::from_utf8(bytes).unwrap();

    // TODO: Fix this unwrap
    // Rust does not support e or E in exponents, so just skip it
    f64::from(i32::from_str_radix(&string[1..], 10).unwrap())
}

fn bytes_to_digit(bytes: &[u8]) -> JSNumber {
    // TODO: Fix this unwrap
    let string = str::from_utf8(bytes).unwrap();

    // TODO: Fix this unwrap
    //i32::from_str_radix(&string, 10).unwrap() as JSNumber
    string.parse::<JSNumber>().unwrap()
}

fn bytes_to_binary(bytes: &[u8]) -> JSNumber {
    // TODO: Fix this unwrap
    let string = str::from_utf8(bytes).unwrap();

    // TODO: Fix this unwrap
    // Rust does not support 0b in binary, so just skip it
    f64::from(i32::from_str_radix(&string[2..], 2).unwrap())
}

fn bytes_to_oct(bytes: &[u8]) -> JSNumber {
    // TODO: Fix this unwrap
    let string = str::from_utf8(bytes).unwrap();

    // TODO: Fix this unwrap
    // Rust does not support 0o in hexadecimal, so just skip it
    f64::from(i32::from_str_radix(&string[2..], 8).unwrap())
}

fn bytes_to_hex(bytes: &[u8]) -> JSNumber {
    // TODO: Fix this unwrap
    let string = str::from_utf8(bytes).unwrap();

    // TODO: Fix this unwrap
    // Rust does not support 0x in hexadecimal, so just skip it
    f64::from(i32::from_str_radix(&string[2..], 16).unwrap())
}

named!(pub numeric_literals<JSNumber>,
    alt!(hex_digits_literal
        | oct_digits_literal
        | binary_digits_literal
        | decimal_digits_literal
        | exponent_part_digits_literal)
);

named!(pub exponent_part_digits_literal<JSNumber>,
    map!(
        recognize!(
            tuple!(
                alt!(tag!("e") | tag!("E")),
                decimal_digits_literal
            )
        ),
        bytes_to_exponent
    )
);

// Support for negating and un-negating numbers will be supplied
// by a function within the language, so we don't parse it
named!(pub decimal_digits_literal<JSNumber>,
    map!(
        digit,
        bytes_to_digit
    )
);

named!(pub binary_digits_literal<JSNumber>,
    map!(
        recognize!(
            tuple!(
                alt!(tag!("0b") | tag!("0B")),
                many1!(binary_digit)
            )
        ),
        bytes_to_binary
    )
);

named!(binary_digit, alt!(tag!("0") | tag!("1")));

named!(pub oct_digits_literal<JSNumber>,
    map!(
        recognize!(
            tuple!(
                alt!(tag!("0o") | tag!("0O")),
                oct_digit,
                opt!(complete!(oct_digit))
            )
        ),
        bytes_to_oct
    )
);

named!(pub hex_digits_literal<JSNumber>,
    map!(
        recognize!(
            tuple!(
                alt!(tag!("0x") | tag!("0X")),
                hex_digit,
                opt!(complete!(hex_digit))
            )
        ),
        bytes_to_hex
    )
);
<commit_msg>Fix rustfmt<commit_after>use nom::{digit, hex_digit, oct_digit};
use std::{str, i32};

// Temporarly use a f64 to represent JS numbers
type JSNumber = f64;

// TODO: Implement proper parsing (and saving) of the number literal.

fn bytes_to_exponent(bytes: &[u8]) -> JSNumber {
    // TODO: Fix this unwrap
    let string = str::from_utf8(bytes).unwrap();

    // TODO: Fix this unwrap
    // Rust does not support e or E in exponents, so just skip it
    f64::from(i32::from_str_radix(&string[1..], 10).unwrap())
}

fn bytes_to_digit(bytes: &[u8]) -> JSNumber {
    // TODO: Fix this unwrap
    let string = str::from_utf8(bytes).unwrap();

    // TODO: Fix this unwrap
    //i32::from_str_radix(&string, 10).unwrap() as JSNumber
    string.parse::<JSNumber>().unwrap()
}

fn bytes_to_binary(bytes: &[u8]) -> JSNumber {
    // TODO: Fix this unwrap
    let string = str::from_utf8(bytes).unwrap();

    // TODO: Fix this unwrap
    // Rust does not support 0b in binary, so just skip it
    f64::from(i32::from_str_radix(&string[2..], 2).unwrap())
}

fn bytes_to_oct(bytes: &[u8]) -> JSNumber {
    // TODO: Fix this unwrap
    let string = str::from_utf8(bytes).unwrap();

    // TODO: Fix this unwrap
    // Rust does not support 0o in hexadecimal, so just skip it
    f64::from(i32::from_str_radix(&string[2..], 8).unwrap())
}

fn bytes_to_hex(bytes: &[u8]) -> JSNumber {
    // TODO: Fix this unwrap
    let string = str::from_utf8(bytes).unwrap();

    // TODO: Fix this unwrap
    // Rust does not support 0x in hexadecimal, so just skip it
    f64::from(i32::from_str_radix(&string[2..], 16).unwrap())
}

named!(pub numeric_literals<JSNumber>,
    alt!(hex_digits_literal
        | oct_digits_literal
        | binary_digits_literal
        | decimal_digits_literal
        | exponent_part_digits_literal)
);

named!(pub exponent_part_digits_literal<JSNumber>,
    map!(
        recognize!(
            tuple!(
                alt!(tag!("e") | tag!("E")),
                decimal_digits_literal
            )
        ),
        bytes_to_exponent
    )
);

// Support for negating and un-negating numbers will be supplied
// by a function within the language, so we don't parse it
named!(pub decimal_digits_literal<JSNumber>,
    map!(
        digit,
        bytes_to_digit
    )
);

named!(pub binary_digits_literal<JSNumber>,
    map!(
        recognize!(
            tuple!(
                alt!(tag!("0b") | tag!("0B")),
                many1!(binary_digit)
            )
        ),
        bytes_to_binary
    )
);

named!(binary_digit, alt!(tag!("0") | tag!("1")));

named!(pub oct_digits_literal<JSNumber>,
    map!(
        recognize!(
            tuple!(
                alt!(tag!("0o") | tag!("0O")),
                oct_digit,
                opt!(complete!(oct_digit))
            )
        ),
        bytes_to_oct
    )
);

named!(pub hex_digits_literal<JSNumber>,
    map!(
        recognize!(
            tuple!(
                alt!(tag!("0x") | tag!("0X")),
                hex_digit,
                opt!(complete!(hex_digit))
            )
        ),
        bytes_to_hex
    )
);
<|endoftext|>


<commit_before>// C preprocessor

use crate::token::*;
use std::cell::RefCell;
use std::collections::HashMap;

thread_local! {
    static MACROS: RefCell<HashMap<String, Macro>> = RefCell::new(HashMap::new());

    static CTX: RefCell<Context> = RefCell::new(Context::new());
}

fn set_ctx(ctx: Context) {
    CTX.with(|c| {
        *c.borrow_mut() = ctx;
    })
}

fn get_ctx() -> Option<Context> {
    CTX.with(|c| {
        Some(c.borrow().clone())
    })
}

fn ctx_pop() {
    CTX.with(|c| {
        let next = c.borrow().next.clone();
        if next.is_some() {
            *c.borrow_mut() = *next.unwrap();
        }
    })
}

fn ctx_output() -> Vec<Token> {
    CTX.with(|c| {
        return c.borrow().output.clone();
    })
}

fn macros_get(key: &String) -> Option<Macro> {
    MACROS.with(|m| {
        match m.borrow().get(key) {
            Some(v) => Some(v.clone()),
            None => None,
        }
    })
}

fn macros_put(key: String, value: Macro) {
    MACROS.with(|m| {
        m.borrow_mut().insert(key, value);
    })
}

#[derive(Clone, Debug)]
struct Context {
    input: Vec<Token>,
    output: Vec<Token>,
    pos: usize,
    next: Option<Box<Context>>,
}

impl Context {
    fn new() -> Context {
        Context {
            input: Vec::new(),
            output: Vec::new(),
            pos: 0,
            next: None,
        }
    }
}

fn new_ctx(next: Option<Context>, input: Vec<Token>) -> Context {
    let mut ctx = Context::new();
    ctx.input = input;
    if next.is_none() {
        ctx.next = None;
    } else {
        ctx.next = Some(Box::new(next.unwrap()));
    }
    return ctx;
}

#[derive(Clone, Debug, PartialEq)]
enum MacroType {
    OBJLIKE,
    FUNCLIKE,
}

#[derive(Clone, Debug)]
struct Macro {
    ty: MacroType,
    tokens: Vec<Token>,
    params: Vec<String>,
}

fn new_macro(ty: MacroType, name: String) {
    let m = Macro {
        ty: ty,
        tokens: Vec::new(),
        params: Vec::new(),
    };
    macros_put(name, m);
}

fn append(v: &mut Vec<Token>) {
    CTX.with(|ctx| {
        ctx.borrow_mut().output.append(v);
    })
}

fn add(t: Token) {
    CTX.with(|ctx| {
        ctx.borrow_mut().output.push(t);
    })
}

fn next() -> Token {
    CTX.with(|ctx| {
        let pos = ctx.borrow().pos;
        assert!(pos < ctx.borrow().input.len());
        ctx.borrow_mut().pos += 1;
        return ctx.borrow().input[pos].clone();
    })
}

fn eof() -> bool {
    CTX.with(|c| {
        let ctx = c.borrow();
        return ctx.pos == ctx.input.len();
    })
}

fn get(ty: TokenType, msg: String) -> Token {
    let t = next();
    if t.ty != ty {
        bad_token(&t, msg);
    }
    return t;
}

fn ident(msg: String) -> String {
    let t = get(TokenType::IDENT, msg);
    return t.name.clone();
}

fn peek() -> Token {
    CTX.with(|c| {
        let ctx = c.borrow();
        return ctx.input[ctx.pos].clone();
    })
}

fn consume(ty: TokenType) -> bool {
    if peek().ty != ty {
        return false;
    }
    CTX.with(|c| {
        c.borrow_mut().pos += 1;
    });
    return true;
}

fn read_until_eol() -> Vec<Token> {
    let mut v = Vec::new();
    while !eof() {
        let t = next();
        if t.ty == TokenType::NEW_LINE {
            break;
        }
        v.push(t.clone());
    }
    return v;
}

fn new_int(val: i32) -> Token {
    let mut t = new_token(TokenType::NUM, 0);
    t.val = val;
    return t;
}

fn new_param(val: i32) -> Token {
    let mut t = new_token(TokenType::PARAM, 0);
    t.val = val;
    return t;
}

fn is_ident(t: &Token, s: &str) -> bool {
    return t.ty == TokenType::IDENT && t.name == s;
}

fn replace_params(m: &mut Macro) {
    let params = m.params.clone();
    let mut tokens = m.tokens.clone();

    // Replace macro parameter tokens with TokenType::PARAM tokens.
    let mut map = HashMap::new();
    for i in 0..params.len() {
        let name = params[i].clone();
        map.insert(name, i as i32);
    }

    for i in 0..tokens.len() {
        let t = tokens[i].clone();
        if t.ty != TokenType::IDENT {
            continue;
        }
        let n = match map.get(&t.name) {
            Some(i) => *i,
            None => {
                continue;
            }
        };
        tokens.remove(i);
        tokens.insert(i, new_param(n));
    }

    // Process '#' followed by a macro parameter.
    let mut v = Vec::new();
    let mut i = 0;
    while i < tokens.len()-1 {
        let t1 = tokens[i].clone();
        let mut t2 = tokens[i+1].clone();

        if t1.ty == TokenType::SHARP && t2.ty == TokenType::PARAM {
            t2.stringize = true;
            v.push(t2);
            i += 1;
        } else {
            v.push(t1);
        }
        i += 1;
    }

    if i == tokens.len() - 1 {
        v.push(tokens[i].clone());
    }
    m.tokens = v;
}

fn read_one_arg() -> Vec<Token> {
    let mut v = Vec::new();
    let start = peek();
    let mut level = 0;

    while !eof() {
        let t = peek();
        if level == 0 {
            if t.ty == TokenType::KET || t.ty == TokenType::COMMA {
                return v;
            }
        }

        next();
        if t.ty == TokenType::BRA {
            level += 1;
        } else if t.ty == TokenType::KET {
            level -= 1;
        }
        v.push(t);
    }
    bad_token(&start, "unclosed macro argument".to_string());
    panic!();
}

fn read_args() -> HashMap<usize, Vec<Token>> {
    let mut v = HashMap::new();
    if consume(TokenType::KET) {
        return v;
    }
    let i = v.len();
    v.insert(i, read_one_arg());
    while !consume(TokenType::KET) {
        get(TokenType::COMMA, "comma expected".to_string());
        let i = v.len();
        v.insert(i, read_one_arg());
    }
    return v;
}

fn stringize(tokens: Vec<Token>) -> Token {
    let mut sb = String::new();

    for i in 0..tokens.len() {
        let t = &tokens[i];
        if i > 0 {
            sb.push(' ');
        }
        sb.push_str(tokstr(&t));
    }
    sb.push('\0');

    let mut t = new_token(TokenType::STR, 0);
    t.len = sb.len();
    t.str_cnt = sb;
    return t;
}

fn add_special_macro(t: & Token) -> bool {
    if is_ident(t, "__LINE__") {
        add(new_int(get_line_number(t)));
        return true;
    }
    return false;
}

fn apply_objlike(m: &mut Macro, _start: & Token) {
    for t in m.tokens.iter() {
        if add_special_macro(t) {
            continue;
        }
        add(t.clone());
    }
}

fn apply_functionlike(m: &mut Macro, start: & Token) {
    get(TokenType::BRA, "comma expected".to_string());
    let args = read_args();
    if m.params.len() != args.len() {
        bad_token(&start, format!("number of parameter does not match ({} != {})", m.params.len(), args.len()).to_string());
        panic!();
    }

    for i in 0..m.tokens.len() {
        let t = &m.tokens[i];
        if add_special_macro(t) {
            continue;
        }

        if t.ty == TokenType::PARAM {
            if t.stringize {
                let j = t.val as usize;
                add(stringize(args.get(&j).unwrap().clone()));
            } else {
                let j = t.val as usize;
                append(&mut args.get(&j).unwrap().clone());
            }
            continue;
        }
        add(t.clone());
    }
}

fn apply(m: &mut Macro, start: & Token) {
    if m.ty == MacroType::OBJLIKE {
        apply_objlike(m, start);
    } else {
        apply_functionlike(m, start);
    }
}

fn funclike_macro(name: String) {
    new_macro(MacroType::FUNCLIKE, name.clone());
    let mut m = macros_get(&name).unwrap();
    m.params.push(ident("parameter name expected".to_string()));
    while !consume(TokenType::KET) {
        get(TokenType::COMMA, "comma expected".to_string());
        m.params.push(ident("parameter name expected".to_string()));
    }
    m.tokens = read_until_eol();

    replace_params(&mut m);

    macros_put(name, m);
}

fn objlike_macro(name: String) {
    new_macro(MacroType::OBJLIKE, name.clone());
    let mut m = macros_get(&name).unwrap();
    m.tokens = read_until_eol();
    macros_put(name, m);
}

fn define() {
    let name = ident("macro name expected".to_string());
    if consume(TokenType::BRA) {
        return funclike_macro(name);
    }
    return objlike_macro(name);
}

fn include() {
    let t = get(TokenType::STR, "string expected".to_string());
    let path = t.str_cnt;
    get(TokenType::NEW_LINE, "newline expected".to_string());
    return append(&mut tokenize(path, false));
}

pub fn preprocess(tokens: Vec<Token>) -> Vec<Token> {
    set_ctx(new_ctx(get_ctx(), tokens.clone()));

    while !eof() {
        let mut t = next();

        if t.ty == TokenType::IDENT {
            match macros_get(&t.name) {
                Some(macro_token) => {
                    apply(&mut macro_token.clone(), &t);
                }
                None => {
                    add(t.clone());
                }
            }
            continue;
        }

        if t.ty != TokenType::SHARP {
            add(t.clone());
            continue;
        }

        t = get(TokenType::IDENT, "identifier expected".to_string());

        if t.name == "define" {
            define();
        } else if t.name == "include" {
            include();
        } else {
            bad_token(&t, "unknown directive".to_string());
        }
    }

    let v = ctx_output();
    ctx_pop();
    return v;
}
<commit_msg>Simplify.<commit_after>// C preprocessor

use crate::token::*;
use std::cell::RefCell;
use std::collections::HashMap;

thread_local! {
    static MACROS: RefCell<HashMap<String, Macro>> = RefCell::new(HashMap::new());

    static CTX: RefCell<Context> = RefCell::new(Context::new());
}

fn set_ctx(ctx: Context) {
    CTX.with(|c| {
        *c.borrow_mut() = ctx;
    })
}

fn get_ctx() -> Option<Context> {
    CTX.with(|c| {
        Some(c.borrow().clone())
    })
}

fn ctx_pop() {
    CTX.with(|c| {
        let next = c.borrow().next.clone();
        if next.is_some() {
            *c.borrow_mut() = *next.unwrap();
        }
    })
}

fn ctx_output() -> Vec<Token> {
    CTX.with(|c| {
        return c.borrow().output.clone();
    })
}

fn macros_get(key: &String) -> Option<Macro> {
    MACROS.with(|m| {
        match m.borrow().get(key) {
            Some(v) => Some(v.clone()),
            None => None,
        }
    })
}

fn macros_put(key: String, value: Macro) {
    MACROS.with(|m| {
        m.borrow_mut().insert(key, value);
    })
}

#[derive(Clone, Debug)]
struct Context {
    input: Vec<Token>,
    output: Vec<Token>,
    pos: usize,
    next: Option<Box<Context>>,
}

impl Context {
    fn new() -> Context {
        Context {
            input: Vec::new(),
            output: Vec::new(),
            pos: 0,
            next: None,
        }
    }
}

fn new_ctx(next: Option<Context>, input: Vec<Token>) -> Context {
    let mut ctx = Context::new();
    ctx.input = input;
    if next.is_none() {
        ctx.next = None;
    } else {
        ctx.next = Some(Box::new(next.unwrap()));
    }
    return ctx;
}

#[derive(Clone, Debug, PartialEq)]
enum MacroType {
    OBJLIKE,
    FUNCLIKE,
}

#[derive(Clone, Debug)]
struct Macro {
    ty: MacroType,
    tokens: Vec<Token>,
    params: Vec<String>,
}

fn new_macro(ty: MacroType, name: String) {
    let m = Macro {
        ty: ty,
        tokens: Vec::new(),
        params: Vec::new(),
    };
    macros_put(name, m);
}

fn append(v: &mut Vec<Token>) {
    CTX.with(|ctx| {
        ctx.borrow_mut().output.append(v);
    })
}

fn add(t: Token) {
    CTX.with(|ctx| {
        ctx.borrow_mut().output.push(t);
    })
}

fn next() -> Token {
    CTX.with(|ctx| {
        let pos = ctx.borrow().pos;
        assert!(pos < ctx.borrow().input.len());
        ctx.borrow_mut().pos += 1;
        return ctx.borrow().input[pos].clone();
    })
}

fn eof() -> bool {
    CTX.with(|c| {
        let ctx = c.borrow();
        return ctx.pos == ctx.input.len();
    })
}

fn get(ty: TokenType, msg: String) -> Token {
    let t = next();
    if t.ty != ty {
        bad_token(&t, msg);
    }
    return t;
}

fn ident(msg: String) -> String {
    let t = get(TokenType::IDENT, msg);
    return t.name.clone();
}

fn peek() -> Token {
    CTX.with(|c| {
        let ctx = c.borrow();
        return ctx.input[ctx.pos].clone();
    })
}

fn consume(ty: TokenType) -> bool {
    if peek().ty != ty {
        return false;
    }
    CTX.with(|c| {
        c.borrow_mut().pos += 1;
    });
    return true;
}

fn read_until_eol() -> Vec<Token> {
    let mut v = Vec::new();
    while !eof() {
        let t = next();
        if t.ty == TokenType::NEW_LINE {
            break;
        }
        v.push(t.clone());
    }
    return v;
}

fn new_int(val: i32) -> Token {
    let mut t = new_token(TokenType::NUM, 0);
    t.val = val;
    return t;
}

fn new_param(val: i32) -> Token {
    let mut t = new_token(TokenType::PARAM, 0);
    t.val = val;
    return t;
}

fn is_ident(t: &Token, s: &str) -> bool {
    return t.ty == TokenType::IDENT && t.name == s;
}

// Replaces macro parameter tokens with TokenType::PARAM tokens.
fn replace_macro_params(m: &mut Macro) {
    let params = m.params.clone();
    let mut tokens = m.tokens.clone();

    let mut map = HashMap::new();
    for i in 0..params.len() {
        let name = params[i].clone();
        map.insert(name, i as i32);
    }

    for i in 0..tokens.len() {
        let t = tokens[i].clone();
        if t.ty != TokenType::IDENT {
            continue;
        }
        let n = match map.get(&t.name) {
            Some(i) => *i,
            None => {
                continue;
            }
        };
        tokens.remove(i);
        tokens.insert(i, new_param(n));
    }
    m.tokens = tokens;
}

// Process '#' followed by a macro parameter.
fn replace_hash_ident(m: &mut Macro) {
    let tokens = m.tokens.clone();
    let mut v = Vec::new();
    let mut i = 0;
    while i < tokens.len()-1 {
        let t1 = tokens[i].clone();
        let mut t2 = tokens[i+1].clone();

        if t1.ty == TokenType::SHARP && t2.ty == TokenType::PARAM {
            t2.stringize = true;
            v.push(t2);
            i += 1;
        } else {
            v.push(t1);
        }
        i += 1;
    }

    if i == tokens.len() - 1 {
        v.push(tokens[i].clone());
    }
    m.tokens = v;
}

fn read_one_arg() -> Vec<Token> {
    let mut v = Vec::new();
    let start = peek();
    let mut level = 0;

    while !eof() {
        let t = peek();
        if level == 0 {
            if t.ty == TokenType::KET || t.ty == TokenType::COMMA {
                return v;
            }
        }

        next();
        if t.ty == TokenType::BRA {
            level += 1;
        } else if t.ty == TokenType::KET {
            level -= 1;
        }
        v.push(t);
    }
    bad_token(&start, "unclosed macro argument".to_string());
    panic!();
}

fn read_args() -> HashMap<usize, Vec<Token>> {
    let mut v = HashMap::new();
    if consume(TokenType::KET) {
        return v;
    }
    let i = v.len();
    v.insert(i, read_one_arg());
    while !consume(TokenType::KET) {
        get(TokenType::COMMA, "comma expected".to_string());
        let i = v.len();
        v.insert(i, read_one_arg());
    }
    return v;
}

fn stringize(tokens: Vec<Token>) -> Token {
    let mut sb = String::new();

    for i in 0..tokens.len() {
        let t = &tokens[i];
        if i > 0 {
            sb.push(' ');
        }
        sb.push_str(tokstr(&t));
    }
    sb.push('\0');

    let mut t = new_token(TokenType::STR, 0);
    t.len = sb.len();
    t.str_cnt = sb;
    return t;
}

fn add_special_macro(t: & Token) -> bool {
    if is_ident(t, "__LINE__") {
        add(new_int(get_line_number(t)));
        return true;
    }
    return false;
}

fn apply_objlike(m: &mut Macro, _start: & Token) {
    for t in m.tokens.iter() {
        if add_special_macro(t) {
            continue;
        }
        add(t.clone());
    }
}

fn apply_functionlike(m: &mut Macro, start: & Token) {
    get(TokenType::BRA, "comma expected".to_string());
    let args = read_args();
    if m.params.len() != args.len() {
        bad_token(&start, format!("number of parameter does not match ({} != {})", m.params.len(), args.len()).to_string());
        panic!();
    }

    for i in 0..m.tokens.len() {
        let t = &m.tokens[i];
        if add_special_macro(t) {
            continue;
        }

        if t.ty == TokenType::PARAM {
            if t.stringize {
                let j = t.val as usize;
                add(stringize(args.get(&j).unwrap().clone()));
            } else {
                let j = t.val as usize;
                append(&mut args.get(&j).unwrap().clone());
            }
            continue;
        }
        add(t.clone());
    }
}

fn apply(m: &mut Macro, start: & Token) {
    if m.ty == MacroType::OBJLIKE {
        apply_objlike(m, start);
    } else {
        apply_functionlike(m, start);
    }
}

fn define_funclike(name: String) {
    new_macro(MacroType::FUNCLIKE, name.clone());
    let mut m = macros_get(&name).unwrap();
    m.params.push(ident("parameter name expected".to_string()));
    while !consume(TokenType::KET) {
        get(TokenType::COMMA, "comma expected".to_string());
        m.params.push(ident("parameter name expected".to_string()));
    }
    m.tokens = read_until_eol();

    replace_macro_params(&mut m);
    replace_hash_ident(&mut m);

    macros_put(name, m);
}

fn define_objlike(name: String) {
    new_macro(MacroType::OBJLIKE, name.clone());
    let mut m = macros_get(&name).unwrap();
    m.tokens = read_until_eol();
    macros_put(name, m);
}

fn define() {
    let name = ident("macro name expected".to_string());
    if consume(TokenType::BRA) {
        return define_funclike(name);
    }
    return define_objlike(name);
}

fn include() {
    let t = get(TokenType::STR, "string expected".to_string());
    let path = t.str_cnt;
    get(TokenType::NEW_LINE, "newline expected".to_string());
    return append(&mut tokenize(path, false));
}

pub fn preprocess(tokens: Vec<Token>) -> Vec<Token> {
    set_ctx(new_ctx(get_ctx(), tokens.clone()));

    while !eof() {
        let mut t = next();

        if t.ty == TokenType::IDENT {
            match macros_get(&t.name) {
                Some(macro_token) => {
                    apply(&mut macro_token.clone(), &t);
                }
                None => {
                    add(t.clone());
                }
            }
            continue;
        }

        if t.ty != TokenType::SHARP {
            add(t.clone());
            continue;
        }

        t = get(TokenType::IDENT, "identifier expected".to_string());

        if t.name == "define" {
            define();
        } else if t.name == "include" {
            include();
        } else {
            bad_token(&t, "unknown directive".to_string());
        }
    }

    let v = ctx_output();
    ctx_pop();
    return v;
}
<|endoftext|>


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

//! This module defines the `__rust_probestack` intrinsic which is used in the
//! implementation of "stack probes" on certain platforms.
//!
//! The purpose of a stack probe is to provide a static guarantee that if a
//! thread has a guard page then a stack overflow is guaranteed to hit that
//! guard page. If a function did not have a stack probe then there's a risk of
//! having a stack frame *larger* than the guard page, so a function call could
//! skip over the guard page entirely and then later hit maybe the heap or
//! another thread, possibly leading to security vulnerabilities such as [The
//! Stack Clash], for example.
//!
//! [The Stack Clash]: https://blog.qualys.com/securitylabs/2017/06/19/the-stack-clash
//!
//! The `__rust_probestack` is called in the prologue of functions whose stack
//! size is larger than the guard page, for example larger than 4096 bytes on
//! x86. This function is then responsible for "touching" all pages relevant to
//! the stack to ensure that that if any of them are the guard page we'll hit
//! them guaranteed.
//!
//! The precise ABI for how this function operates is defined by LLVM. There's
//! no real documentation as to what this is, so you'd basically need to read
//! the LLVM source code for reference. Often though the test cases can be
//! illuminating as to the ABI that's generated, or just looking at the output
//! of `llc`.
//!
//! Note that `#[naked]` is typically used here for the stack probe because the
//! ABI corresponds to no actual ABI.
//!
//! Finally it's worth noting that at the time of this writing LLVM only has
//! support for stack probes on x86 and x86_64. There's no support for stack
//! probes on any other architecture like ARM or PowerPC64. LLVM I'm sure would
//! be more than welcome to accept such a change!

#![cfg(not(windows))] // Windows already has builtins to do this

#[naked]
#[no_mangle]
#[cfg(target_arch = "x86_64")]
pub unsafe extern fn __rust_probestack() {
    // Our goal here is to touch each page between %rsp+8 and %rsp+8-%rax,
    // ensuring that if any pages are unmapped we'll make a page fault.
    //
    // The ABI here is that the stack frame size is located in `%eax`. Upon
    // return we're not supposed to modify `%esp` or `%eax`.
    asm!("
        lea    8(%rsp),%r11     // rsp before calling this routine -> r11

        // Main loop, taken in one page increments. We're decrementing r11 by
        // a page each time until there's less than a page remaining. We're
        // guaranteed that this function isn't called unless there's more than a
        // page needed
    2:
        sub    $$0x1000,%r11
        test   %r11,(%r11)
        sub    $$0x1000,%rax
        cmp    $$0x1000,%rax
        ja     2b

        // Finish up the last remaining stack space requested, getting the last
        // bits out of rax
        sub    %rax,%r11
        test   %r11,(%r11)

        // We now know that %r11 is (%rsp + 8 - %rax) so to recover rax
        // we calculate (%rsp + 8) - %r11 which will give us %rax
        lea    8(%rsp),%rax
        sub    %r11,%rax

        ret
    ");
    ::core::intrinsics::unreachable();
}

#[naked]
#[no_mangle]
#[cfg(target_arch = "x86")]
pub unsafe extern fn __rust_probestack() {
    // This is the same as x86_64 above, only translated for 32-bit sizes. Note
    // that on Unix we're expected to restore everything as it was, this
    // function basically can't tamper with anything.
    //
    // The ABI here is the same as x86_64, except everything is 32-bits large.
    asm!("
        push   %ecx
        lea    8(%esp),%ecx
    2:
        sub    $$0x1000,%ecx
        test   %ecx,(%ecx)
        sub    $$0x1000,%eax
        cmp    $$0x1000,%eax
        ja     2b

        sub    %eax,%ecx
        test   %ecx,(%ecx)

        lea    8(%esp),%eax
        sub    %ecx,%eax
        pop    %ecx
        ret
    ");
    ::core::intrinsics::unreachable();
}
<commit_msg>Tweak definition of probestack functions<commit_after>// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! This module defines the `__rust_probestack` intrinsic which is used in the
//! implementation of "stack probes" on certain platforms.
//!
//! The purpose of a stack probe is to provide a static guarantee that if a
//! thread has a guard page then a stack overflow is guaranteed to hit that
//! guard page. If a function did not have a stack probe then there's a risk of
//! having a stack frame *larger* than the guard page, so a function call could
//! skip over the guard page entirely and then later hit maybe the heap or
//! another thread, possibly leading to security vulnerabilities such as [The
//! Stack Clash], for example.
//!
//! [The Stack Clash]: https://blog.qualys.com/securitylabs/2017/06/19/the-stack-clash
//!
//! The `__rust_probestack` is called in the prologue of functions whose stack
//! size is larger than the guard page, for example larger than 4096 bytes on
//! x86. This function is then responsible for "touching" all pages relevant to
//! the stack to ensure that that if any of them are the guard page we'll hit
//! them guaranteed.
//!
//! The precise ABI for how this function operates is defined by LLVM. There's
//! no real documentation as to what this is, so you'd basically need to read
//! the LLVM source code for reference. Often though the test cases can be
//! illuminating as to the ABI that's generated, or just looking at the output
//! of `llc`.
//!
//! Note that `#[naked]` is typically used here for the stack probe because the
//! ABI corresponds to no actual ABI.
//!
//! Finally it's worth noting that at the time of this writing LLVM only has
//! support for stack probes on x86 and x86_64. There's no support for stack
//! probes on any other architecture like ARM or PowerPC64. LLVM I'm sure would
//! be more than welcome to accept such a change!

#![cfg(not(windows))] // Windows already has builtins to do this

#[naked]
#[no_mangle]
#[cfg(target_arch = "x86_64")]
pub unsafe extern fn __rust_probestack() {
    // Our goal here is to touch each page between %rsp+8 and %rsp+8-%rax,
    // ensuring that if any pages are unmapped we'll make a page fault.
    //
    // The ABI here is that the stack frame size is located in `%eax`. Upon
    // return we're not supposed to modify `%esp` or `%eax`.
    asm!("
        mov    %rax,%r11        // duplicate %rax as we're clobbering %r11

        // Main loop, taken in one page increments. We're decrementing rsp by
        // a page each time until there's less than a page remaining. We're
        // guaranteed that this function isn't called unless there's more than a
        // page needed.
        //
        // Note that we're also testing against `8(%rsp)` to account for the 8
        // bytes pushed on the stack orginally with our return address. Using
        // `8(%rsp)` simulates us testing the stack pointer in the caller's
        // context.
    2:
        sub    $$0x1000,%rsp
        test   %rsp,8(%rsp)
        sub    $$0x1000,%r11
        cmp    $$0x1000,%r11
        ja     2b

        // Finish up the last remaining stack space requested, getting the last
        // bits out of r11
        sub    %r11,%rsp
        test   %rsp,8(%rsp)

        // Restore the stack pointer to what it previously was when entering
        // this function. The caller will readjust the stack pointer after we
        // return.
        add    %rax,%rsp

        ret
    ");
    ::core::intrinsics::unreachable();
}

#[naked]
#[no_mangle]
#[cfg(target_arch = "x86")]
pub unsafe extern fn __rust_probestack() {
    // This is the same as x86_64 above, only translated for 32-bit sizes. Note
    // that on Unix we're expected to restore everything as it was, this
    // function basically can't tamper with anything.
    //
    // The ABI here is the same as x86_64, except everything is 32-bits large.
    asm!("
        push   %ecx
        mov    %eax,%ecx
    2:
        sub    $$0x1000,%esp
        test   %esp,8(%esp)
        sub    $$0x1000,%ecx
        cmp    $$0x1000,%ecx
        ja     2b

        sub    %ecx,%esp
        test   %esp,8(%esp)

        add    %eax,%esp
        pop    %ecx
        ret
    ");
    ::core::intrinsics::unreachable();
}
<|endoftext|>