<commit_before>// Copyright 2014 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.

// ignore-pretty
// compile-flags:--test
// exec-env:RUST_TEST_TASKS=1

// Tests for the connect_timeout() function on a TcpStream. This runs with only
// one test task to ensure that errors are timeouts, not file descriptor
// exhaustion.

#![feature(macro_rules, globs)]
#![allow(experimental)]

extern crate native;
extern crate green;
extern crate rustuv;

#[cfg(test)] #[start]
fn start(argc: int, argv: **u8) -> int {
    green::start(argc, argv, rustuv::event_loop, __test::main)
}

macro_rules! iotest (
    { fn $name:ident() $b:block $($a:attr)* } => (
        mod $name {
            #![allow(unused_imports)]

            use std::io::*;
            use std::io::net::tcp::*;
            use std::io::test::*;
            use std::io;

            fn f() $b

            $($a)* #[test] fn green() { f() }
            $($a)* #[test] fn native() {
                use native;
                let (tx, rx) = channel();
                native::task::spawn(proc() { tx.send(f()) });
                rx.recv();
            }
        }
    )
)

iotest!(fn eventual_timeout() {
    use native;
    let addr = next_test_ip4();

    // Use a native task to receive connections because it turns out libuv is
    // really good at accepting connections and will likely run out of file
    // descriptors before timing out.
    let (tx1, rx1) = channel();
    let (_tx2, rx2) = channel::<()>();
    native::task::spawn(proc() {
        let _l = TcpListener::bind(addr).unwrap().listen();
        tx1.send(());
        let _ = rx2.recv_opt();
    });
    rx1.recv();

    let mut v = Vec::new();
    for _ in range(0, 10000) {
        match TcpStream::connect_timeout(addr, 100) {
            Ok(e) => v.push(e),
            Err(ref e) if e.kind == io::TimedOut => return,
            Err(e) => fail!("other error: {}", e),
        }
    }
    fail!("never timed out!");
})

iotest!(fn timeout_success() {
    let addr = next_test_ip4();
    let _l = TcpListener::bind(addr).unwrap().listen();

    assert!(TcpStream::connect_timeout(addr, 1000).is_ok());
})

iotest!(fn timeout_error() {
    let addr = next_test_ip4();

    assert!(TcpStream::connect_timeout(addr, 1000).is_err());
})
<commit_msg>test: Ignore tcp-connect-timeout on freebsd<commit_after>// Copyright 2014 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.

// ignore-pretty
// compile-flags:--test
// exec-env:RUST_TEST_TASKS=1

// Tests for the connect_timeout() function on a TcpStream. This runs with only
// one test task to ensure that errors are timeouts, not file descriptor
// exhaustion.

#![feature(macro_rules, globs)]
#![allow(experimental)]

extern crate native;
extern crate green;
extern crate rustuv;

#[cfg(test)] #[start]
fn start(argc: int, argv: **u8) -> int {
    green::start(argc, argv, rustuv::event_loop, __test::main)
}

macro_rules! iotest (
    { fn $name:ident() $b:block $(#[$a:meta])* } => (
        mod $name {
            #![allow(unused_imports)]

            use std::io::*;
            use std::io::net::tcp::*;
            use std::io::test::*;
            use std::io;

            fn f() $b

            $(#[$a])* #[test] fn green() { f() }
            $(#[$a])* #[test] fn native() {
                use native;
                let (tx, rx) = channel();
                native::task::spawn(proc() { tx.send(f()) });
                rx.recv();
            }
        }
    )
)

iotest!(fn eventual_timeout() {
    use native;
    let addr = next_test_ip4();

    // Use a native task to receive connections because it turns out libuv is
    // really good at accepting connections and will likely run out of file
    // descriptors before timing out.
    let (tx1, rx1) = channel();
    let (_tx2, rx2) = channel::<()>();
    native::task::spawn(proc() {
        let _l = TcpListener::bind(addr).unwrap().listen();
        tx1.send(());
        let _ = rx2.recv_opt();
    });
    rx1.recv();

    let mut v = Vec::new();
    for _ in range(0, 10000) {
        match TcpStream::connect_timeout(addr, 100) {
            Ok(e) => v.push(e),
            Err(ref e) if e.kind == io::TimedOut => return,
            Err(e) => fail!("other error: {}", e),
        }
    }
    fail!("never timed out!");
} #[ignore(cfg(target_os = "freebsd"))])

iotest!(fn timeout_success() {
    let addr = next_test_ip4();
    let _l = TcpListener::bind(addr).unwrap().listen();

    assert!(TcpStream::connect_timeout(addr, 1000).is_ok());
})

iotest!(fn timeout_error() {
    let addr = next_test_ip4();

    assert!(TcpStream::connect_timeout(addr, 1000).is_err());
})
<|endoftext|>


<commit_before><commit_msg>Use single argument to vec4 constructor<commit_after><|endoftext|>


<commit_before>
use super::*;

/// Error structure used to represent the COM `IErrorInfo`.
pub struct ComError {
    pub description : String,
    pub hresult : HRESULT,
}

#[cfg(windows)]
#[link(name = "oleaut32")]
extern "system" {
    pub fn SetErrorInfo(
        dw_reserved: u32,
        errorinfo: RawComPtr,
    ) -> HRESULT;
}

#[cfg(not(windows))]
pub fn SetErrorInfo(
    dw_reserved: u32,
    errorinfo: RawComPtr,
) -> HRESULT { S_OK }

/// Error info COM object data.
struct ErrorInfo { pub description : String }

/// `CoClass` implementation for the `ErrorInfo` COM object.
impl CoClass for ErrorInfo {

    /// Virtual table type.
    type VTableList = &'static IErrorInfoVtbl;

    /// Virtual table constructor.
    fn create_vtable_list() -> Self::VTableList {
        &IErrorInfoVtbl {
            __base : IUnknownVtbl {
                query_interface : ComBox::< Self >::query_interface_ptr,
                add_ref : ComBox::< Self >::add_ref_ptr,
                release : ComBox::< Self >::release_ptr,
            },
            get_guid : get_guid_impl,
            get_source : get_source_impl,
            get_description : get_description_impl,
            get_help_file : get_help_file_impl,
            get_help_context : get_help_context_impl,
        }
    }

    /// Query interface implementation. Supports only `IUnknown` and `IErrorInfo`.
    fn query_interface(
        vtables : &Self::VTableList,
        riid : REFIID,
    ) -> ComResult< RawComPtr >
    {
        if riid.is_null() { return Err( E_NOINTERFACE ) }
        unsafe { match *riid {
            super::IID_IUnknown | super::IID_IErrorInfo =>
                Ok( vtables as *const _ as RawComPtr ),
            _ => Err( E_NOINTERFACE ),
        } }
    }

    /// `IErrorInfo` itself doesn't support error info.
    ///
    /// We don't want to override errors when doing error processing.
    fn interface_supports_error_info( _riid : REFIID ) -> bool { false }
}

/// `IErrorInfo` virtual table.
struct IErrorInfoVtbl {
    pub __base: IUnknownVtbl,
    pub get_guid: unsafe extern "stdcall" fn(
            RawComPtr, *mut GUID ) -> HRESULT,
    pub get_source: unsafe extern "stdcall" fn(
            RawComPtr, *mut BStr ) -> HRESULT,
    pub get_description: unsafe extern "stdcall" fn(
            RawComPtr, *mut BStr ) -> HRESULT,
    pub get_help_file: unsafe extern "stdcall" fn(
            RawComPtr, *mut BStr ) -> HRESULT,
    pub get_help_context: unsafe extern "stdcall" fn(
            RawComPtr, *mut u32 ) -> HRESULT
}

/// `IErrorInfo::GetGuid` raw COM implementation.
unsafe extern "stdcall" fn get_guid_impl(
    _this : RawComPtr,
    guid : *mut GUID
) -> HRESULT
{
    if guid.is_null() {
        return E_INVALIDARG;
    }

    *guid = GUID::zero_guid();
    S_OK
}

/// `IErrorInfo::GetSource` raw COM implementation.
unsafe extern "stdcall" fn get_source_impl(
    _this : RawComPtr,
    source : *mut BStr
) -> HRESULT
{
    if source.is_null() {
        return E_INVALIDARG;
    }

    *source = BStr::string_to_bstr( "Intercom-Component" );
    S_OK
}

/// `IErrorInfo::GetDescription` raw COM implementation.
unsafe extern "stdcall" fn get_description_impl(
    this : RawComPtr,
    desc : *mut BStr
) -> HRESULT
{
    if desc.is_null() {
        return E_INVALIDARG;
    }

    let cb = ComBox::< ErrorInfo >::from_ptr( this );
    *desc = BStr::string_to_bstr( &cb.description );
    S_OK
}

/// `IErrorInfo::GetHelpFile` raw COM implementation.
unsafe extern "stdcall" fn get_help_file_impl(
    _this : RawComPtr,
    help_file : *mut BStr
) -> HRESULT
{
    if help_file.is_null() {
        return E_INVALIDARG;
    }

    *help_file = BStr::string_to_bstr( "" );
    S_OK
}

/// `IErrorInfo::GetHelpContext` raw COM implementation.
unsafe extern "stdcall" fn get_help_context_impl(
    _this : RawComPtr,
    help_context : *mut u32
) -> HRESULT
{
    if help_context.is_null() {
        return E_INVALIDARG;
    }

    *help_context = 0;
    S_OK
}

/// Extracts the HRESULT from the error result and stores the extended error
/// information in thread memory so it can be fetched by the COM client.
pub fn return_hresult< E >( error : E ) -> HRESULT
    where E : Into< ComError >
{
    // Convet the error.
    let com_error = error.into();

    // Construct the COM class used for IErrorInfo. The class contains the
    // description in memory.
    let mut info = ComBox::< ErrorInfo >::new( ErrorInfo {
        description: com_error.description
    } );

    // Get the IErrorInfo interface and set it in thread memory.
    let mut error_ptr : RawComPtr = std::ptr::null_mut();
    unsafe {

        // We are intentionally ignoring the HRESULT codes here. We don't
        // want to override the original error HRESULT with these codes.
        ComBox::query_interface(
                info.as_mut(),
                &IID_IErrorInfo,
                &mut error_ptr );
        SetErrorInfo( 0, error_ptr );

        // SetErrorInfo took ownership of the error.
        // Forget it from the Box.
        Box::into_raw( info );
    }

    // Return the HRESULT of the original error.
    com_error.hresult
}

/// Gets the last COM error that occurred on the current thread.
pub fn get_last_error< E >() -> E
    where E : From< ComError >
{
    panic!( "Not implemented" );
}
<commit_msg>More Linux lint fixes<commit_after>
use super::*;

/// Error structure used to represent the COM `IErrorInfo`.
pub struct ComError {
    pub description : String,
    pub hresult : HRESULT,
}

#[cfg(windows)]
#[link(name = "oleaut32")]
extern "system" {
    pub fn SetErrorInfo(
        dw_reserved: u32,
        errorinfo: RawComPtr,
    ) -> HRESULT;
}

#[cfg(not(windows))]
#[allow(non_snake_case)]
pub fn SetErrorInfo(
    _dw_reserved: u32,
    _errorinfo: RawComPtr,
) -> HRESULT { S_OK }

/// Error info COM object data.
struct ErrorInfo { pub description : String }

/// `CoClass` implementation for the `ErrorInfo` COM object.
impl CoClass for ErrorInfo {

    /// Virtual table type.
    type VTableList = &'static IErrorInfoVtbl;

    /// Virtual table constructor.
    fn create_vtable_list() -> Self::VTableList {
        &IErrorInfoVtbl {
            __base : IUnknownVtbl {
                query_interface : ComBox::< Self >::query_interface_ptr,
                add_ref : ComBox::< Self >::add_ref_ptr,
                release : ComBox::< Self >::release_ptr,
            },
            get_guid : get_guid_impl,
            get_source : get_source_impl,
            get_description : get_description_impl,
            get_help_file : get_help_file_impl,
            get_help_context : get_help_context_impl,
        }
    }

    /// Query interface implementation. Supports only `IUnknown` and `IErrorInfo`.
    fn query_interface(
        vtables : &Self::VTableList,
        riid : REFIID,
    ) -> ComResult< RawComPtr >
    {
        if riid.is_null() { return Err( E_NOINTERFACE ) }
        unsafe { match *riid {
            super::IID_IUnknown | super::IID_IErrorInfo =>
                Ok( vtables as *const _ as RawComPtr ),
            _ => Err( E_NOINTERFACE ),
        } }
    }

    /// `IErrorInfo` itself doesn't support error info.
    ///
    /// We don't want to override errors when doing error processing.
    fn interface_supports_error_info( _riid : REFIID ) -> bool { false }
}

/// `IErrorInfo` virtual table.
struct IErrorInfoVtbl {
    pub __base: IUnknownVtbl,
    pub get_guid: unsafe extern "stdcall" fn(
            RawComPtr, *mut GUID ) -> HRESULT,
    pub get_source: unsafe extern "stdcall" fn(
            RawComPtr, *mut BStr ) -> HRESULT,
    pub get_description: unsafe extern "stdcall" fn(
            RawComPtr, *mut BStr ) -> HRESULT,
    pub get_help_file: unsafe extern "stdcall" fn(
            RawComPtr, *mut BStr ) -> HRESULT,
    pub get_help_context: unsafe extern "stdcall" fn(
            RawComPtr, *mut u32 ) -> HRESULT
}

/// `IErrorInfo::GetGuid` raw COM implementation.
unsafe extern "stdcall" fn get_guid_impl(
    _this : RawComPtr,
    guid : *mut GUID
) -> HRESULT
{
    if guid.is_null() {
        return E_INVALIDARG;
    }

    *guid = GUID::zero_guid();
    S_OK
}

/// `IErrorInfo::GetSource` raw COM implementation.
unsafe extern "stdcall" fn get_source_impl(
    _this : RawComPtr,
    source : *mut BStr
) -> HRESULT
{
    if source.is_null() {
        return E_INVALIDARG;
    }

    *source = BStr::string_to_bstr( "Intercom-Component" );
    S_OK
}

/// `IErrorInfo::GetDescription` raw COM implementation.
unsafe extern "stdcall" fn get_description_impl(
    this : RawComPtr,
    desc : *mut BStr
) -> HRESULT
{
    if desc.is_null() {
        return E_INVALIDARG;
    }

    let cb = ComBox::< ErrorInfo >::from_ptr( this );
    *desc = BStr::string_to_bstr( &cb.description );
    S_OK
}

/// `IErrorInfo::GetHelpFile` raw COM implementation.
unsafe extern "stdcall" fn get_help_file_impl(
    _this : RawComPtr,
    help_file : *mut BStr
) -> HRESULT
{
    if help_file.is_null() {
        return E_INVALIDARG;
    }

    *help_file = BStr::string_to_bstr( "" );
    S_OK
}

/// `IErrorInfo::GetHelpContext` raw COM implementation.
unsafe extern "stdcall" fn get_help_context_impl(
    _this : RawComPtr,
    help_context : *mut u32
) -> HRESULT
{
    if help_context.is_null() {
        return E_INVALIDARG;
    }

    *help_context = 0;
    S_OK
}

/// Extracts the HRESULT from the error result and stores the extended error
/// information in thread memory so it can be fetched by the COM client.
pub fn return_hresult< E >( error : E ) -> HRESULT
    where E : Into< ComError >
{
    // Convet the error.
    let com_error = error.into();

    // Construct the COM class used for IErrorInfo. The class contains the
    // description in memory.
    let mut info = ComBox::< ErrorInfo >::new( ErrorInfo {
        description: com_error.description
    } );

    // Get the IErrorInfo interface and set it in thread memory.
    let mut error_ptr : RawComPtr = std::ptr::null_mut();
    unsafe {

        // We are intentionally ignoring the HRESULT codes here. We don't
        // want to override the original error HRESULT with these codes.
        ComBox::query_interface(
                info.as_mut(),
                &IID_IErrorInfo,
                &mut error_ptr );
        SetErrorInfo( 0, error_ptr );

        // SetErrorInfo took ownership of the error.
        // Forget it from the Box.
        Box::into_raw( info );
    }

    // Return the HRESULT of the original error.
    com_error.hresult
}

/// Gets the last COM error that occurred on the current thread.
pub fn get_last_error< E >() -> E
    where E : From< ComError >
{
    panic!( "Not implemented" );
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate modparse_derive;
#[macro_use]
extern crate derive_builder;
#[macro_use]
extern crate torchrs;
#[macro_use]
extern crate clap;

use torchrs::autograd::{Variable, VariableArgs, VarAccess};
use torchrs::optim;
use torchrs::optim::OptimVal;
use torchrs::nn;
use torchrs::nn::{InitModuleStruct, GetFieldStruct, ModIntf, ModDelegate, Module};
use torchrs::nn::functional as F;
use torchrs::utils::data as D;
use torchrs::utils::torchvision::{datasets, transforms};
use torchrs::tensor;

use clap::{Arg, App};
#[derive(Clone, Builder)]
struct NetArgs {
    #[builder(default="64")]
    batch_size: usize,
    #[builder(default="1000")]
    test_batch_size: usize,
    #[builder(default="10")]
    epochs: u32,
    #[builder(default="0.01")]
    lr: f32,
    #[builder(default="0.5")]
    momentum: f32,
    #[builder(default="false")]
    cuda: bool,
    #[builder(default="1")]
    seed: usize,
    #[builder(default="10")]
    log_interval: usize,
}

impl Default for NetArgs {
    fn default() -> Self {
        NetArgsBuilder::default().build().unwrap()
    }
}

fn parse_args() -> NetArgs {
    let mut args = NetArgs::default();
    let matches = App::new("Torch.rs MNIST Example")
        .version("0.1")
        .author("Some Guy")
        .about("recognize handwritten digits")
        .arg(Arg::with_name("batch-size").takes_value(true))
        .arg(Arg::with_name("test-batch-size").takes_value(true))
        .arg(Arg::with_name("epochs").takes_value(true))
        .arg(Arg::with_name("lr").takes_value(true))
        .arg(Arg::with_name("momentum").takes_value(true))
        .arg(Arg::with_name("no-cuda").takes_value(false))
        .arg(Arg::with_name("seed").takes_value(true))
        .arg(Arg::with_name("log-interval").takes_value(true))
        .get_matches();
    let (b_size, tb_size, epochs, lr, momentum, seed, log_int) =
        (value_t!(matches.value_of("batch-size"), usize).ok(),
         value_t!(matches.value_of("test-batch-size"), usize).ok(),
         value_t!(matches.value_of("epochs"), u32).ok(),
         value_t!(matches.value_of("lr"), f32).ok(),
         value_t!(matches.value_of("momentum"), f32).ok(),
         value_t!(matches.value_of("seed"), usize).ok(),
         value_t!(matches.value_of("log-interval"), usize).ok());
    if let Some(b_size) = b_size {
        args.batch_size = b_size
    }
    if let Some(tb_size) = tb_size {
        args.test_batch_size = tb_size
    }
    if let Some(epochs) = epochs {
        args.epochs = epochs
    }
    if let Some(lr) = lr {
        args.lr = lr
    }
    if let Some(momentum) = momentum {
        args.momentum = momentum
    }
    if let Some(seed) = seed {
        args.seed = seed
    }
    if let Some(log_int) = log_int {
        args.log_interval = log_int
    }
    //XXX CUDA?
    args
}

#[derive(ModParse)]
struct Net<T: tensor::NumLimits> {
    delegate: nn::Module<T>,
    conv1: nn::Conv2d<T>,
    conv2: nn::Conv2d<T>,
    fc1: nn::Linear<T>,
    fc2: nn::Linear<T>,
}

impl<T: tensor::NumLimits> Net<T> {
    pub fn new() -> Net<T> {
        Net {
                delegate: nn::Module::new(),
                conv1: nn::Conv2d::build(1, 10, (5, 5)).done(),
                conv2: nn::Conv2d::build(10, 20, (5, 5)).done(),
                fc1: nn::Linear::build(320, 50).done(),
                fc2: nn::Linear::build(50, 10).done(),
            }
            .init_module()
    }
}
impl_mod_delegate!(Net);

impl<T: tensor::NumLimits> ModIntf<T> for Net<T> {
    fn forward(&mut self, args: &mut Variable<T>) -> Variable<T> {
        let pool_val = F::MaxPool2dArgs::default();
        let mut dropout_val = F::DropoutArgs::default();
        dropout_val.training = self.delegate.training;
        let mut x = F::relu(F::max_pool2d(self.conv1.f(args.clone()), (2, 2), &pool_val));
        x = F::relu(F::max_pool2d(F::dropout2d(self.conv2.f(x), &dropout_val),
                                  (2, 2),
                                  &pool_val));
        x = x.view([-1, 320]);
        x = F::relu(self.fc1.f(x));
        x = F::dropout(x, &dropout_val);
        x = self.fc2.f(x);
        F::log_softmax(x)
    }
}

fn train(model: &mut Net<f32>,
         args: &NetArgs,
         train_loader: &D::BatchLoader<f32, i64>,
         epoch: u32,
         optimizer: &mut optim::OptIntf<f32>) {
    model.train();
    for (batch_idx, (ref data, ref target)) in train_loader.iter().enumerate() {
        let (mut data, target) = if args.cuda {
            (Variable::new(data.cuda(None).clone()), Variable::new(target.cuda(None).clone()))
        } else {
            (Variable::new(data.clone()), Variable::new(target.clone()))
        };
        optimizer.zero_grad(model);
        let output = model.f(data.clone());
        let mut loss = F::nll_loss(output, target, &F::NLLLossArgs::default());
        loss.backward();
        optimizer.step(model);
        model.free_graph();
        if batch_idx % args.log_interval == 0 {
            println!("Train Epoch: {} [{}/{} ({:.0}%)]\tLoss: {:.6}",
                     epoch,
                     batch_idx * data.data().len(),
                     train_loader.dataset.len(),
                     100. * (batch_idx as f32) / train_loader.len() as f32,
                     loss.data()[0]);
        }
    }
}

fn test(model: &mut Net<f32>, args: &NetArgs, test_loader: &D::BatchLoader<f32, i64>) {
    model.eval();
    let mut test_loss = 0.;
    let mut correct = 0;
    for (ref data, ref target) in test_loader.iter() {
        let varargs = VariableArgs {
            volatile: true,
            ..Default::default()
        };
        let (data, mut target) = if args.cuda {
            (Variable::new_args(data.cuda(None).clone(), &varargs),
             Variable::new(target.cuda(None).clone()))
        } else {
            (Variable::new_args(data.clone(), &varargs), Variable::new(target.clone()))
        };
        let mut output = model.f(data);
        test_loss += F::nll_loss(output.clone(), target.clone(), &F::NLLLossArgs::default());
        let pred = output.data().max_reduce(1).1;
        correct += pred.eq_tensor(&*target.data()).cpu().sum() as u32;
    }
    test_loss /= test_loader.len() as f32;
    println!("\nTest set: Average loss: {:.4}, Accuracy: {}/{} ({:.0}%)\n",
             test_loss,
             correct,
             test_loader.dataset.len(),
             100. * (correct as f32) / (test_loader.dataset.len() as f32))

}

fn main() {
    let args = parse_args();
    let train_loader: D::BatchLoader<f32, i64> = D::DataLoader::build()
        .batch_size(args.batch_size)
        .done(datasets::MNIST::<f32>::build("../data")
                  .download(false)
                  .done(None));

    let test_loader: D::BatchLoader<f32, i64> = D::DataLoader::build()
        .batch_size(args.batch_size)
        .done(datasets::MNIST::<f32>::build("../data")
                  .train(false)
                  .done(None));
    let mut model = Net::new();
    let mut optimizer = optim::SGD::new(map_opt!{"lr" => args.lr, "momentum" => args.momentum});
    for epoch in 1..args.epochs + 1 {
        train(&mut model, &args, &train_loader, epoch, &mut optimizer);
        test(&mut model, &args, &test_loader);
    }
}
<commit_msg>fix compile<commit_after>#[macro_use]
extern crate modparse_derive;
#[macro_use]
extern crate derive_builder;
#[macro_use]
extern crate torchrs;
#[macro_use]
extern crate clap;

use torchrs::autograd::{Variable, VariableArgs, VarAccess};
use torchrs::optim;
use torchrs::optim::OptimVal;
use torchrs::nn;
use torchrs::nn::{InitModuleStruct, GetFieldStruct, ModIntf, ModDelegate, Module};
use torchrs::nn::functional as F;
use torchrs::utils::data as D;
use torchrs::utils::torchvision::{datasets, transforms};
use torchrs::tensor;

use clap::{Arg, App};
#[derive(Clone, Builder)]
struct NetArgs {
    #[builder(default="64")]
    batch_size: usize,
    #[builder(default="1000")]
    test_batch_size: usize,
    #[builder(default="10")]
    epochs: u32,
    #[builder(default="0.01")]
    lr: f32,
    #[builder(default="0.5")]
    momentum: f32,
    #[builder(default="false")]
    cuda: bool,
    #[builder(default="1")]
    seed: usize,
    #[builder(default="10")]
    log_interval: usize,
}

impl Default for NetArgs {
    fn default() -> Self {
        NetArgsBuilder::default().build().unwrap()
    }
}

fn parse_args() -> NetArgs {
    let mut args = NetArgs::default();
    let matches = App::new("Torch.rs MNIST Example")
        .version("0.1")
        .author("Some Guy")
        .about("recognize handwritten digits")
        .arg(Arg::with_name("batch-size").takes_value(true))
        .arg(Arg::with_name("test-batch-size").takes_value(true))
        .arg(Arg::with_name("epochs").takes_value(true))
        .arg(Arg::with_name("lr").takes_value(true))
        .arg(Arg::with_name("momentum").takes_value(true))
        .arg(Arg::with_name("no-cuda").takes_value(false))
        .arg(Arg::with_name("seed").takes_value(true))
        .arg(Arg::with_name("log-interval").takes_value(true))
        .get_matches();
    let (b_size, tb_size, epochs, lr, momentum, seed, log_int) =
        (value_t!(matches.value_of("batch-size"), usize).ok(),
         value_t!(matches.value_of("test-batch-size"), usize).ok(),
         value_t!(matches.value_of("epochs"), u32).ok(),
         value_t!(matches.value_of("lr"), f32).ok(),
         value_t!(matches.value_of("momentum"), f32).ok(),
         value_t!(matches.value_of("seed"), usize).ok(),
         value_t!(matches.value_of("log-interval"), usize).ok());
    if let Some(b_size) = b_size {
        args.batch_size = b_size
    }
    if let Some(tb_size) = tb_size {
        args.test_batch_size = tb_size
    }
    if let Some(epochs) = epochs {
        args.epochs = epochs
    }
    if let Some(lr) = lr {
        args.lr = lr
    }
    if let Some(momentum) = momentum {
        args.momentum = momentum
    }
    if let Some(seed) = seed {
        args.seed = seed
    }
    if let Some(log_int) = log_int {
        args.log_interval = log_int
    }
    //XXX CUDA?
    args
}

#[derive(ModParse)]
struct Net<T: tensor::NumLimits> {
    delegate: nn::Module<T>,
    conv1: nn::Conv2d<T>,
    conv2: nn::Conv2d<T>,
    fc1: nn::Linear<T>,
    fc2: nn::Linear<T>,
}

impl<T: tensor::NumLimits> Net<T> {
    pub fn new() -> Net<T> {
        Net {
                delegate: nn::Module::new(),
                conv1: nn::Conv2d::build(1, 10, (5, 5)).done(),
                conv2: nn::Conv2d::build(10, 20, (5, 5)).done(),
                fc1: nn::Linear::build(320, 50).done(),
                fc2: nn::Linear::build(50, 10).done(),
            }
            .init_module()
    }
}
impl_mod_delegate!(Net);

impl<T: tensor::NumLimits> ModIntf<T> for Net<T> {
    fn forward(&mut self, args: &mut Variable<T>) -> Variable<T> {
        let pool_val = F::MaxPool2dArgs::default();
        let mut dropout_val = F::DropoutArgs::default();
        dropout_val.training = self.delegate.training;
        let mut x = F::relu(F::max_pool2d(self.conv1.f(args.clone()), (2, 2), &pool_val));
        x = F::relu(F::max_pool2d(F::dropout2d(self.conv2.f(x), &dropout_val),
                                  (2, 2),
                                  &pool_val));
        x = x.view([-1, 320]);
        x = F::relu(self.fc1.f(x));
        x = F::dropout(x, &dropout_val);
        x = self.fc2.f(x);
        F::log_softmax(x)
    }
}

fn train(model: &mut Net<f32>,
         args: &NetArgs,
         train_loader: &D::BatchLoader<f32, i64>,
         epoch: u32,
         optimizer: &mut optim::OptIntf<f32>) {
    model.train();
    for (batch_idx, (ref data, ref target)) in train_loader.iter().enumerate() {
        let (mut data, target) = if args.cuda {
            (Variable::new(data.cuda(None).clone()), Variable::new(target.cuda(None).clone()))
        } else {
            (Variable::new(data.clone()), Variable::new(target.clone()))
        };
        optimizer.zero_grad(model);
        let output = model.f(data.clone());
        let mut loss = F::nll_loss(output, target, &F::NLLLossArgs::default());
        loss.backward();
        optimizer.step(model);
        model.free_graph();
        if batch_idx % args.log_interval == 0 {
            println!("Train Epoch: {} [{}/{} ({:.0}%)]\tLoss: {:.6}",
                     epoch,
                     batch_idx * data.data().len(),
                     train_loader.dataset.len(),
                     100. * (batch_idx as f32) / train_loader.len() as f32,
                     loss.data()[0]);
        }
    }
}

fn test(model: &mut Net<f32>, args: &NetArgs, test_loader: &D::BatchLoader<f32, i64>) {
    model.eval();
    let mut test_loss = 0.;
    let mut correct = 0;
    for (ref data, ref target) in test_loader.iter() {
        let varargs = VariableArgs {
            volatile: true,
            ..Default::default()
        };
        let (data, mut target) = if args.cuda {
            (Variable::new_args(data.cuda(None).clone(), &varargs),
             Variable::new(target.cuda(None).clone()))
        } else {
            (Variable::new_args(data.clone(), &varargs), Variable::new(target.clone()))
        };
        let mut output = model.f(data);
        test_loss += F::nll_loss(output.clone(), target.clone(), &F::NLLLossArgs::default());
        let pred = output.data().max_reduce(1).1;
        correct += pred.eq_tensor(&*target.data()).cpu().sum::<u32>();
    }
    test_loss /= test_loader.len() as f32;
    println!("\nTest set: Average loss: {:.4}, Accuracy: {}/{} ({:.0}%)\n",
             test_loss,
             correct,
             test_loader.dataset.len(),
             100. * (correct as f32) / (test_loader.dataset.len() as f32))

}

fn main() {
    let args = parse_args();
    let train_loader: D::BatchLoader<f32, i64> = D::DataLoader::build()
        .batch_size(args.batch_size)
        .done(datasets::MNIST::<f32>::build("../data")
                  .download(false)
                  .done(None));

    let test_loader: D::BatchLoader<f32, i64> = D::DataLoader::build()
        .batch_size(args.batch_size)
        .done(datasets::MNIST::<f32>::build("../data")
                  .train(false)
                  .done(None));
    let mut model = Net::new();
    let mut optimizer = optim::SGD::new(map_opt!{"lr" => args.lr, "momentum" => args.momentum});
    for epoch in 1..args.epochs + 1 {
        train(&mut model, &args, &train_loader, epoch, &mut optimizer);
        test(&mut model, &args, &test_loader);
    }
}
<|endoftext|>


<commit_before>/* This example asks the rtkit service to make our thread realtime priority.
   Rtkit puts a few limitations on us to let us become realtime, such as setting
   RLIMIT_RTTIME correctly, hence the syscalls. */

#![feature(old_orphan_check, libc)]

extern crate "dbus-rs" as dbus;
extern crate libc;

use std::cmp;
use libc::c_int;

use dbus::{Connection, BusType, Error, Props, MessageItem, Message};

/* External C stuff. Currently only working for x86_64 */

#[allow(non_upper_case_globals)]
#[cfg(target_arch = "x86_64")]
static SYS_gettid: c_int = 186;

static RLIMIT_RTTIME: c_int = 15;

#[repr(C)]
#[cfg(target_arch = "x86_64")]
#[derive(Copy)]
struct rlimit {
    rlim_cur: u64,
    rlim_max: u64,
}

extern "C" {
    fn syscall(num: c_int, ...) -> c_int;

    fn getrlimit(resource: c_int, rlim: *mut rlimit) -> c_int;
    fn setrlimit(resource: c_int, rlim: *const rlimit) -> c_int;
}

/* Error handling */

struct MyError(String);

impl<T: std::fmt::Debug> std::error::FromError<T> for MyError {
    fn from_error(err: T) -> MyError {
        MyError(format!("{:?}", err))
    }
}

fn item_as_i64(i: MessageItem) -> Result<i64, String> {
    match i {
        MessageItem::Int32(i) => Ok(i as i64),
        MessageItem::Int64(i) => Ok(i),
        _ => Err(format!("Property is not integer ({:?})", i))
    }
}

fn rtkit_set_realtime(c: &Connection, thread: u64, prio: u32) -> Result<(), Error> {
    let mut m = Message::new_method_call("org.freedesktop.RealtimeKit1", "/org/freedesktop/RealtimeKit1",
        "org.freedesktop.RealtimeKit1", "MakeThreadRealtime").unwrap();
    m.append_items(&[MessageItem::UInt64(thread), MessageItem::UInt32(prio)]);
    let mut r = try!(c.send_with_reply_and_block(m, 10000));
    r.as_result().map(|_| ())
}

fn make_realtime(prio: u32) -> Result<u32, MyError> {
    let c = try!(Connection::get_private(BusType::System));

    let p = Props::new(&c, "org.freedesktop.RealtimeKit1", "/org/freedesktop/RealtimeKit1",
        "org.freedesktop.RealtimeKit1", 10000);

    // Make sure we don't fail by wanting too much
    let max_prio = try!(item_as_i64(try!(p.get("MaxRealtimePriority")))) as u32;
    let prio = cmp::min(prio, max_prio);

    // Enforce RLIMIT_RTPRIO, also a must before asking rtkit for rtprio
    let max_rttime = try!(item_as_i64(try!(p.get("RTTimeNSecMax")))) as u64;
    let new_limit = rlimit { rlim_cur: max_rttime, rlim_max: max_rttime };
    let mut old_limit = new_limit;
    if unsafe { getrlimit(RLIMIT_RTTIME, &mut old_limit) } < 0 {
        return Err(MyError("getrlimit failed".to_string()));
    }
    if unsafe { setrlimit(RLIMIT_RTTIME, &new_limit) } < 0 {
        return Err(MyError("setrlimit failed".to_string()));
    }

    // Finally, let's ask rtkit to make us realtime
    let thread_id = unsafe { syscall(SYS_gettid) };
    let r = rtkit_set_realtime(&c, thread_id as u64, prio);

    if r.is_err() {
        unsafe { setrlimit(RLIMIT_RTTIME, &old_limit) };
    }

    try!(r);
    Ok(prio)
}

fn main() {
    match make_realtime(5) {
        Ok(n) => println!("Got rtprio, level {}", n),
        Err(e) => println!("No rtprio: {}", e.0),
    }
}
<commit_msg>Remove useless old_orphan_check feature<commit_after>/* This example asks the rtkit service to make our thread realtime priority.
   Rtkit puts a few limitations on us to let us become realtime, such as setting
   RLIMIT_RTTIME correctly, hence the syscalls. */

#![feature(libc)]

extern crate "dbus-rs" as dbus;
extern crate libc;

use std::cmp;
use libc::c_int;

use dbus::{Connection, BusType, Error, Props, MessageItem, Message};

/* External C stuff. Currently only working for x86_64 */

#[allow(non_upper_case_globals)]
#[cfg(target_arch = "x86_64")]
static SYS_gettid: c_int = 186;

static RLIMIT_RTTIME: c_int = 15;

#[repr(C)]
#[cfg(target_arch = "x86_64")]
#[derive(Copy)]
struct rlimit {
    rlim_cur: u64,
    rlim_max: u64,
}

extern "C" {
    fn syscall(num: c_int, ...) -> c_int;

    fn getrlimit(resource: c_int, rlim: *mut rlimit) -> c_int;
    fn setrlimit(resource: c_int, rlim: *const rlimit) -> c_int;
}

/* Error handling */

struct MyError(String);

impl<T: std::fmt::Debug> std::error::FromError<T> for MyError {
    fn from_error(err: T) -> MyError {
        MyError(format!("{:?}", err))
    }
}

fn item_as_i64(i: MessageItem) -> Result<i64, String> {
    match i {
        MessageItem::Int32(i) => Ok(i as i64),
        MessageItem::Int64(i) => Ok(i),
        _ => Err(format!("Property is not integer ({:?})", i))
    }
}

fn rtkit_set_realtime(c: &Connection, thread: u64, prio: u32) -> Result<(), Error> {
    let mut m = Message::new_method_call("org.freedesktop.RealtimeKit1", "/org/freedesktop/RealtimeKit1",
        "org.freedesktop.RealtimeKit1", "MakeThreadRealtime").unwrap();
    m.append_items(&[MessageItem::UInt64(thread), MessageItem::UInt32(prio)]);
    let mut r = try!(c.send_with_reply_and_block(m, 10000));
    r.as_result().map(|_| ())
}

fn make_realtime(prio: u32) -> Result<u32, MyError> {
    let c = try!(Connection::get_private(BusType::System));

    let p = Props::new(&c, "org.freedesktop.RealtimeKit1", "/org/freedesktop/RealtimeKit1",
        "org.freedesktop.RealtimeKit1", 10000);

    // Make sure we don't fail by wanting too much
    let max_prio = try!(item_as_i64(try!(p.get("MaxRealtimePriority")))) as u32;
    let prio = cmp::min(prio, max_prio);

    // Enforce RLIMIT_RTPRIO, also a must before asking rtkit for rtprio
    let max_rttime = try!(item_as_i64(try!(p.get("RTTimeNSecMax")))) as u64;
    let new_limit = rlimit { rlim_cur: max_rttime, rlim_max: max_rttime };
    let mut old_limit = new_limit;
    if unsafe { getrlimit(RLIMIT_RTTIME, &mut old_limit) } < 0 {
        return Err(MyError("getrlimit failed".to_string()));
    }
    if unsafe { setrlimit(RLIMIT_RTTIME, &new_limit) } < 0 {
        return Err(MyError("setrlimit failed".to_string()));
    }

    // Finally, let's ask rtkit to make us realtime
    let thread_id = unsafe { syscall(SYS_gettid) };
    let r = rtkit_set_realtime(&c, thread_id as u64, prio);

    if r.is_err() {
        unsafe { setrlimit(RLIMIT_RTTIME, &old_limit) };
    }

    try!(r);
    Ok(prio)
}

fn main() {
    match make_realtime(5) {
        Ok(n) => println!("Got rtprio, level {}", n),
        Err(e) => println!("No rtprio: {}", e.0),
    }
}
<|endoftext|>


<commit_before>/// This function is where the kernel sets up IRQ handlers
/// It is increcibly unsafe, and should be minimal in nature
/// It must create the IDT with the correct entries, those entries are
/// defined in other files inside of the `arch` module

use core::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, AtomicUsize, ATOMIC_USIZE_INIT, Ordering};

use acpi;
use allocator;
use device;
use externs::memset;
use gdt;
use idt;
use interrupt;
use memory::{self, Frame};
use paging::{self, entry, Page, PhysicalAddress, VirtualAddress};

/// Test of zero values in BSS.
static BSS_TEST_ZERO: usize = 0;
/// Test of non-zero values in data.
static DATA_TEST_NONZERO: usize = 0xFFFFFFFFFFFFFFFF;
/// Test of zero values in thread BSS
#[thread_local]
static mut TBSS_TEST_ZERO: usize = 0;
/// Test of non-zero values in thread data.
#[thread_local]
static mut TDATA_TEST_NONZERO: usize = 0xFFFFFFFFFFFFFFFF;

pub static AP_READY: AtomicBool = ATOMIC_BOOL_INIT;
static BSP_READY: AtomicBool = ATOMIC_BOOL_INIT;
static HEAP_TABLE: AtomicUsize = ATOMIC_USIZE_INIT;

extern {
    /// Kernel main function
    fn kmain() -> !;
    /// Kernel main for APs
    fn kmain_ap(id: usize) -> !;
}

/// The entry to Rust, all things must be initialized
#[no_mangle]
pub unsafe extern fn kstart() -> ! {
    {
        extern {
            /// The starting byte of the _.bss_ (uninitialized data) segment.
            static mut __bss_start: u8;
            /// The ending byte of the _.bss_ (uninitialized data) segment.
            static mut __bss_end: u8;
            /// The end of the kernel
            static mut __end: u8;
        }

        // Zero BSS, this initializes statics that are set to 0
        {
            let start_ptr = &mut __bss_start as *mut u8;
            let end_ptr = & __bss_end as *const u8 as usize;

            if start_ptr as usize <= end_ptr {
                let size = end_ptr - start_ptr as usize;
                memset(start_ptr, 0, size);
            }

            assert_eq!(BSS_TEST_ZERO, 0);
            assert_eq!(DATA_TEST_NONZERO, 0xFFFFFFFFFFFFFFFF);
        }

        // Initialize memory management
        memory::init(0, &__end as *const u8 as usize);

        // TODO: allocate a stack
        let stack_start = 0x00080000;
        let stack_end = 0x0009F000;

        // Initialize paging
        let (mut active_table, tcb_offset) = paging::init(0, stack_start, stack_end);

        // Set up GDT
        gdt::init(tcb_offset, stack_end);

        // Set up IDT
        idt::init();

        // Test tdata and tbss
        {
            assert_eq!(TBSS_TEST_ZERO, 0);
            TBSS_TEST_ZERO += 1;
            assert_eq!(TBSS_TEST_ZERO, 1);
            assert_eq!(TDATA_TEST_NONZERO, 0xFFFFFFFFFFFFFFFF);
            TDATA_TEST_NONZERO -= 1;
            assert_eq!(TDATA_TEST_NONZERO, 0xFFFFFFFFFFFFFFFE);
        }

        // Reset AP variables
        AP_READY.store(false, Ordering::SeqCst);
        BSP_READY.store(false, Ordering::SeqCst);
        HEAP_TABLE.store(0, Ordering::SeqCst);

        // Map heap
        {
            // Map heap pages
            let heap_start_page = Page::containing_address(VirtualAddress::new(::KERNEL_HEAP_OFFSET));
            let heap_end_page = Page::containing_address(VirtualAddress::new(::KERNEL_HEAP_OFFSET + ::KERNEL_HEAP_SIZE-1));
            for page in Page::range_inclusive(heap_start_page, heap_end_page) {
                active_table.map(page, entry::WRITABLE | entry::NO_EXECUTE);
            }

            // Init the allocator
            allocator::init(::KERNEL_HEAP_OFFSET, ::KERNEL_HEAP_SIZE);

            // Send heap page table to APs
            let index = heap_start_page.p4_index();

            let p4 = active_table.p4();
            {
                let entry = &p4[index];
                if let Some(frame) = entry.pointed_frame() {
                    HEAP_TABLE.store(frame.start_address().get(), Ordering::SeqCst);
                } else {
                    panic!("heap does not have PML4 entry");
                }
            }
        }

        // Initialize devices
        device::init(&mut active_table);

        // Read ACPI tables, starts APs
        acpi::init(&mut active_table);

        BSP_READY.store(true, Ordering::SeqCst);
    }

    kmain();
}

/// Entry to rust for an AP
pub unsafe extern fn kstart_ap(cpu_id: usize, page_table: usize, stack_start: usize, stack_end: usize) -> ! {
    {
        assert_eq!(BSS_TEST_ZERO, 0);
        assert_eq!(DATA_TEST_NONZERO, 0xFFFFFFFFFFFFFFFF);

        // Initialize paging
        let (mut active_table, tcb_offset) = paging::init(cpu_id, stack_start, stack_end);

        // Set up GDT for AP
        gdt::init(tcb_offset, stack_end);

        // Set up IDT for AP
        idt::init();

        // Test tdata and tbss
        {
            assert_eq!(TBSS_TEST_ZERO, 0);
            TBSS_TEST_ZERO += 1;
            assert_eq!(TBSS_TEST_ZERO, 1);
            assert_eq!(TDATA_TEST_NONZERO, 0xFFFFFFFFFFFFFFFF);
            TDATA_TEST_NONZERO -= 1;
            assert_eq!(TDATA_TEST_NONZERO, 0xFFFFFFFFFFFFFFFE);
        }

        // Copy heap PML4
        {
            while HEAP_TABLE.load(Ordering::SeqCst) == 0 {
                interrupt::pause();
            }
            let frame = Frame::containing_address(PhysicalAddress::new(HEAP_TABLE.load(Ordering::SeqCst)));

            let page = Page::containing_address(VirtualAddress::new(::KERNEL_HEAP_OFFSET));
            let p4 = active_table.p4_mut();
            let entry = &mut p4[page.p4_index()];
            assert!(entry.is_unused());
            entry.set(frame, entry::PRESENT | entry::WRITABLE);
        }

        // Init devices for AP
        device::init_ap(&mut active_table);

        AP_READY.store(true, Ordering::SeqCst);
    }

    while ! BSP_READY.load(Ordering::SeqCst) {
        interrupt::pause();
    }

    kmain_ap(cpu_id);
}

pub unsafe fn usermode(ip: usize, sp: usize) -> ! {
    // Go to usermode
    asm!("mov rax, 0x2B # Set segment pointers
        mov ds, ax
        mov es, ax
        mov fs, ax
        mov gs, ax

        push rax # Push stack segment
        push rbx # Push stack pointer
        mov rax, 3 << 12 | 1 << 9 # Set IOPL and interrupt enable flag
        push rax # Push rflags
        mov rax, 0x23
        push rax # Push code segment
        push rcx # Push rip
        iretq"
        :
        : "{rbx}"(sp), "{rcx}"(ip)
        : "rax", "sp"
        : "intel", "volatile");
    unreachable!();
}
<commit_msg>Cleanup usermode jump<commit_after>/// This function is where the kernel sets up IRQ handlers
/// It is increcibly unsafe, and should be minimal in nature
/// It must create the IDT with the correct entries, those entries are
/// defined in other files inside of the `arch` module

use core::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, AtomicUsize, ATOMIC_USIZE_INIT, Ordering};

use acpi;
use allocator;
use device;
use externs::memset;
use gdt;
use idt;
use interrupt;
use memory::{self, Frame};
use paging::{self, entry, Page, PhysicalAddress, VirtualAddress};

/// Test of zero values in BSS.
static BSS_TEST_ZERO: usize = 0;
/// Test of non-zero values in data.
static DATA_TEST_NONZERO: usize = 0xFFFFFFFFFFFFFFFF;
/// Test of zero values in thread BSS
#[thread_local]
static mut TBSS_TEST_ZERO: usize = 0;
/// Test of non-zero values in thread data.
#[thread_local]
static mut TDATA_TEST_NONZERO: usize = 0xFFFFFFFFFFFFFFFF;

pub static AP_READY: AtomicBool = ATOMIC_BOOL_INIT;
static BSP_READY: AtomicBool = ATOMIC_BOOL_INIT;
static HEAP_TABLE: AtomicUsize = ATOMIC_USIZE_INIT;

extern {
    /// Kernel main function
    fn kmain() -> !;
    /// Kernel main for APs
    fn kmain_ap(id: usize) -> !;
}

/// The entry to Rust, all things must be initialized
#[no_mangle]
pub unsafe extern fn kstart() -> ! {
    {
        extern {
            /// The starting byte of the _.bss_ (uninitialized data) segment.
            static mut __bss_start: u8;
            /// The ending byte of the _.bss_ (uninitialized data) segment.
            static mut __bss_end: u8;
            /// The end of the kernel
            static mut __end: u8;
        }

        // Zero BSS, this initializes statics that are set to 0
        {
            let start_ptr = &mut __bss_start as *mut u8;
            let end_ptr = & __bss_end as *const u8 as usize;

            if start_ptr as usize <= end_ptr {
                let size = end_ptr - start_ptr as usize;
                memset(start_ptr, 0, size);
            }

            assert_eq!(BSS_TEST_ZERO, 0);
            assert_eq!(DATA_TEST_NONZERO, 0xFFFFFFFFFFFFFFFF);
        }

        // Initialize memory management
        memory::init(0, &__end as *const u8 as usize);

        // TODO: allocate a stack
        let stack_start = 0x00080000;
        let stack_end = 0x0009F000;

        // Initialize paging
        let (mut active_table, tcb_offset) = paging::init(0, stack_start, stack_end);

        // Set up GDT
        gdt::init(tcb_offset, stack_end);

        // Set up IDT
        idt::init();

        // Test tdata and tbss
        {
            assert_eq!(TBSS_TEST_ZERO, 0);
            TBSS_TEST_ZERO += 1;
            assert_eq!(TBSS_TEST_ZERO, 1);
            assert_eq!(TDATA_TEST_NONZERO, 0xFFFFFFFFFFFFFFFF);
            TDATA_TEST_NONZERO -= 1;
            assert_eq!(TDATA_TEST_NONZERO, 0xFFFFFFFFFFFFFFFE);
        }

        // Reset AP variables
        AP_READY.store(false, Ordering::SeqCst);
        BSP_READY.store(false, Ordering::SeqCst);
        HEAP_TABLE.store(0, Ordering::SeqCst);

        // Map heap
        {
            // Map heap pages
            let heap_start_page = Page::containing_address(VirtualAddress::new(::KERNEL_HEAP_OFFSET));
            let heap_end_page = Page::containing_address(VirtualAddress::new(::KERNEL_HEAP_OFFSET + ::KERNEL_HEAP_SIZE-1));
            for page in Page::range_inclusive(heap_start_page, heap_end_page) {
                active_table.map(page, entry::WRITABLE | entry::NO_EXECUTE);
            }

            // Init the allocator
            allocator::init(::KERNEL_HEAP_OFFSET, ::KERNEL_HEAP_SIZE);

            // Send heap page table to APs
            let index = heap_start_page.p4_index();

            let p4 = active_table.p4();
            {
                let entry = &p4[index];
                if let Some(frame) = entry.pointed_frame() {
                    HEAP_TABLE.store(frame.start_address().get(), Ordering::SeqCst);
                } else {
                    panic!("heap does not have PML4 entry");
                }
            }
        }

        // Initialize devices
        device::init(&mut active_table);

        // Read ACPI tables, starts APs
        acpi::init(&mut active_table);

        BSP_READY.store(true, Ordering::SeqCst);
    }

    kmain();
}

/// Entry to rust for an AP
pub unsafe extern fn kstart_ap(cpu_id: usize, page_table: usize, stack_start: usize, stack_end: usize) -> ! {
    {
        assert_eq!(BSS_TEST_ZERO, 0);
        assert_eq!(DATA_TEST_NONZERO, 0xFFFFFFFFFFFFFFFF);

        // Initialize paging
        let (mut active_table, tcb_offset) = paging::init(cpu_id, stack_start, stack_end);

        // Set up GDT for AP
        gdt::init(tcb_offset, stack_end);

        // Set up IDT for AP
        idt::init();

        // Test tdata and tbss
        {
            assert_eq!(TBSS_TEST_ZERO, 0);
            TBSS_TEST_ZERO += 1;
            assert_eq!(TBSS_TEST_ZERO, 1);
            assert_eq!(TDATA_TEST_NONZERO, 0xFFFFFFFFFFFFFFFF);
            TDATA_TEST_NONZERO -= 1;
            assert_eq!(TDATA_TEST_NONZERO, 0xFFFFFFFFFFFFFFFE);
        }

        // Copy heap PML4
        {
            while HEAP_TABLE.load(Ordering::SeqCst) == 0 {
                interrupt::pause();
            }
            let frame = Frame::containing_address(PhysicalAddress::new(HEAP_TABLE.load(Ordering::SeqCst)));

            let page = Page::containing_address(VirtualAddress::new(::KERNEL_HEAP_OFFSET));
            let p4 = active_table.p4_mut();
            let entry = &mut p4[page.p4_index()];
            assert!(entry.is_unused());
            entry.set(frame, entry::PRESENT | entry::WRITABLE);
        }

        // Init devices for AP
        device::init_ap(&mut active_table);

        AP_READY.store(true, Ordering::SeqCst);
    }

    while ! BSP_READY.load(Ordering::SeqCst) {
        interrupt::pause();
    }

    kmain_ap(cpu_id);
}

pub unsafe fn usermode(ip: usize, sp: usize) -> ! {
    // Go to usermode
    asm!("xchg bx, bx
        mov ds, ax
        mov es, ax
        mov fs, ax
        mov gs, ax
        push rax
        push rbx
        push rcx
        push rdx
        push rsi
        iretq"
        : // No output because it never returns
        :   "{rax}"(gdt::GDT_USER_DATA << 3 | 3), // Stack segment
            "{rbx}"(sp), // Stack pointer
            "{rcx}"(3 << 12 | 1 << 9), // Flags - Set IOPL and interrupt enable flag
            "{rdx}"(gdt::GDT_USER_CODE << 3 | 3), // Code segment
            "{rsi}"(ip) // IP
        : // No clobers because it never returns
        : "intel", "volatile");
    unreachable!();
}
<|endoftext|>


<commit_before>//! Tests for custom cargo commands and other global command features.

use std::env;
use std::fs;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::str;

use cargo_test_support::registry::Package;
use cargo_test_support::tools::echo_subcommand;
use cargo_test_support::{
    basic_bin_manifest, cargo_exe, cargo_process, paths, project, project_in_home,
};

fn path() -> Vec<PathBuf> {
    env::split_paths(&env::var_os("PATH").unwrap_or_default()).collect()
}

#[cargo_test]
fn list_commands_with_descriptions() {
    let p = project().build();
    p.cargo("--list")
        .with_stdout_contains(
            "    build                Compile a local package and all of its dependencies",
        )
        // Assert that `read-manifest` prints the right one-line description followed by another
        // command, indented.
        .with_stdout_contains(
            "    read-manifest        Print a JSON representation of a Cargo.toml manifest.",
        )
        .run();
}

#[cargo_test]
fn list_builtin_aliases_with_descriptions() {
    let p = project().build();
    p.cargo("--list")
        .with_stdout_contains("    b                    alias: build")
        .with_stdout_contains("    c                    alias: check")
        .with_stdout_contains("    r                    alias: run")
        .with_stdout_contains("    t                    alias: test")
        .run();
}

#[cargo_test]
fn list_custom_aliases_with_descriptions() {
    let p = project_in_home("proj")
        .file(
            &paths::home().join(".cargo").join("config"),
            r#"
            [alias]
            myaliasstr = "foo --bar"
            myaliasvec = ["foo", "--bar"]
        "#,
        )
        .build();

    p.cargo("--list")
        .with_stdout_contains("    myaliasstr           alias: foo --bar")
        .with_stdout_contains("    myaliasvec           alias: foo --bar")
        .run();
}

#[cargo_test]
fn list_dedupe() {
    let p = project()
        .executable(Path::new("path-test-1").join("cargo-dupe"), "")
        .executable(Path::new("path-test-2").join("cargo-dupe"), "")
        .build();

    let mut path = path();
    path.push(p.root().join("path-test-1"));
    path.push(p.root().join("path-test-2"));
    let path = env::join_paths(path.iter()).unwrap();

    p.cargo("--list")
        .env("PATH", &path)
        .with_stdout_contains_n("    dupe", 1)
        .run();
}

#[cargo_test]
fn list_command_looks_at_path() {
    let proj = project()
        .executable(Path::new("path-test").join("cargo-1"), "")
        .build();

    let mut path = path();
    path.push(proj.root().join("path-test"));
    let path = env::join_paths(path.iter()).unwrap();
    let output = cargo_process("-v --list")
        .env("PATH", &path)
        .exec_with_output()
        .unwrap();
    let output = str::from_utf8(&output.stdout).unwrap();
    assert!(
        output.contains("\n    1                   "),
        "missing 1: {}",
        output
    );
}

#[cargo_test]
fn list_command_handles_known_external_commands() {
    let p = project()
        .executable(Path::new("path-test").join("cargo-fmt"), "")
        .build();

    let fmt_desc = "    fmt                  Formats all bin and lib files of the current crate using rustfmt.";

    // Without path - fmt isn't there
    p.cargo("--list")
        .env("PATH", "")
        .with_stdout_does_not_contain(fmt_desc)
        .run();

    // With path - fmt is there with known description
    let mut path = path();
    path.push(p.root().join("path-test"));
    let path = env::join_paths(path.iter()).unwrap();

    p.cargo("--list")
        .env("PATH", &path)
        .with_stdout_contains(fmt_desc)
        .run();
}

#[cargo_test]
fn list_command_resolves_symlinks() {
    let proj = project()
        .symlink(cargo_exe(), Path::new("path-test").join("cargo-2"))
        .build();

    let mut path = path();
    path.push(proj.root().join("path-test"));
    let path = env::join_paths(path.iter()).unwrap();
    let output = cargo_process("-v --list")
        .env("PATH", &path)
        .exec_with_output()
        .unwrap();
    let output = str::from_utf8(&output.stdout).unwrap();
    assert!(
        output.contains("\n    2                   "),
        "missing 2: {}",
        output
    );
}

#[cargo_test]
fn find_closest_capital_c_to_c() {
    cargo_process("C")
        .with_status(101)
        .with_stderr_contains(
            "\
error: no such subcommand: `C`

<tab>Did you mean `c`?
",
        )
        .run();
}

#[cargo_test]
fn find_closest_capital_b_to_b() {
    cargo_process("B")
        .with_status(101)
        .with_stderr_contains(
            "\
error: no such subcommand: `B`

<tab>Did you mean `b`?
",
        )
        .run();
}

#[cargo_test]
fn find_closest_biuld_to_build() {
    cargo_process("biuld")
        .with_status(101)
        .with_stderr_contains(
            "\
error: no such subcommand: `biuld`

<tab>Did you mean `build`?
",
        )
        .run();

    // But, if we actually have `biuld`, it must work!
    // https://github.com/rust-lang/cargo/issues/5201
    Package::new("cargo-biuld", "1.0.0")
        .file(
            "src/main.rs",
            r#"
                fn main() {
                    println!("Similar, but not identical to, build");
                }
            "#,
        )
        .publish();

    cargo_process("install cargo-biuld").run();
    cargo_process("biuld")
        .with_stdout("Similar, but not identical to, build\n")
        .run();
    cargo_process("--list")
        .with_stdout_contains(
            "    build                Compile a local package and all of its dependencies\n",
        )
        .with_stdout_contains("    biuld\n")
        .run();
}

#[cargo_test]
fn find_closest_alias() {
    let root = paths::root();
    let my_home = root.join("my_home");
    fs::create_dir(&my_home).unwrap();
    fs::write(
        &my_home.join("config"),
        r#"
            [alias]
            myalias = "build"
        "#,
    )
    .unwrap();

    cargo_process("myalais")
        .env("CARGO_HOME", &my_home)
        .with_status(101)
        .with_stderr_contains(
            "\
error: no such subcommand: `myalais`

<tab>Did you mean `myalias`?
",
        )
        .run();

    // But, if no alias is defined, it must not suggest one!
    cargo_process("myalais")
        .with_status(101)
        .with_stderr_contains(
            "\
error: no such subcommand: `myalais`
",
        )
        .with_stderr_does_not_contain(
            "\
<tab>Did you mean `myalias`?
",
        )
        .run();
}

// If a subcommand is more than an edit distance of 3 away, we don't make a suggestion.
#[cargo_test]
fn find_closest_dont_correct_nonsense() {
    cargo_process("there-is-no-way-that-there-is-a-command-close-to-this")
        .cwd(&paths::root())
        .with_status(101)
        .with_stderr(
            "[ERROR] no such subcommand: `there-is-no-way-that-there-is-a-command-close-to-this`\n
<tab>View all installed commands with `cargo --list`",
        )
        .run();
}

#[cargo_test]
fn displays_subcommand_on_error() {
    cargo_process("invalid-command")
        .with_status(101)
        .with_stderr(
            "[ERROR] no such subcommand: `invalid-command`\n
<tab>View all installed commands with `cargo --list`",
        )
        .run();
}

#[cargo_test]
fn override_cargo_home() {
    let root = paths::root();
    let my_home = root.join("my_home");
    fs::create_dir(&my_home).unwrap();
    fs::write(
        &my_home.join("config"),
        r#"
            [cargo-new]
            vcs = "none"
        "#,
    )
    .unwrap();

    cargo_process("new foo").env("CARGO_HOME", &my_home).run();

    assert!(!paths::root().join("foo/.git").is_dir());

    cargo_process("new foo2").run();

    assert!(paths::root().join("foo2/.git").is_dir());
}

#[cargo_test]
fn cargo_subcommand_env() {
    let src = format!(
        r#"
        use std::env;

        fn main() {{
            println!("{{}}", env::var("{}").unwrap());
        }}
        "#,
        cargo::CARGO_ENV
    );

    let p = project()
        .at("cargo-envtest")
        .file("Cargo.toml", &basic_bin_manifest("cargo-envtest"))
        .file("src/main.rs", &src)
        .build();

    let target_dir = p.target_debug_dir();

    p.cargo("build").run();
    assert!(p.bin("cargo-envtest").is_file());

    let cargo = cargo_exe().canonicalize().unwrap();
    let mut path = path();
    path.push(target_dir);
    let path = env::join_paths(path.iter()).unwrap();

    cargo_process("envtest")
        .env("PATH", &path)
        .with_stdout(cargo.to_str().unwrap())
        .run();
}

#[cargo_test]
fn cargo_subcommand_args() {
    let p = echo_subcommand();
    let cargo_foo_bin = p.bin("cargo-echo");
    assert!(cargo_foo_bin.is_file());

    let mut path = path();
    path.push(p.target_debug_dir());
    let path = env::join_paths(path.iter()).unwrap();

    cargo_process("echo bar -v --help")
        .env("PATH", &path)
        .with_stdout("echo bar -v --help")
        .run();
}

#[cargo_test]
fn explain() {
    cargo_process("--explain E0001")
        .with_stdout_contains(
            "This error suggests that the expression arm corresponding to the noted pattern",
        )
        .run();
}

#[cargo_test]
fn closed_output_ok() {
    // Checks that closed output doesn't cause an error.
    let mut p = cargo_process("--list").build_command();
    p.stdout(Stdio::piped()).stderr(Stdio::piped());
    let mut child = p.spawn().unwrap();
    // Close stdout
    drop(child.stdout.take());
    // Read stderr
    let mut s = String::new();
    child
        .stderr
        .as_mut()
        .unwrap()
        .read_to_string(&mut s)
        .unwrap();
    let status = child.wait().unwrap();
    assert!(status.success());
    assert!(s.is_empty(), "{}", s);
}
<commit_msg>stop mixing of newlines and `\n` in tests<commit_after>//! Tests for custom cargo commands and other global command features.

use std::env;
use std::fs;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::str;

use cargo_test_support::registry::Package;
use cargo_test_support::tools::echo_subcommand;
use cargo_test_support::{
    basic_bin_manifest, cargo_exe, cargo_process, paths, project, project_in_home,
};

fn path() -> Vec<PathBuf> {
    env::split_paths(&env::var_os("PATH").unwrap_or_default()).collect()
}

#[cargo_test]
fn list_commands_with_descriptions() {
    let p = project().build();
    p.cargo("--list")
        .with_stdout_contains(
            "    build                Compile a local package and all of its dependencies",
        )
        // Assert that `read-manifest` prints the right one-line description followed by another
        // command, indented.
        .with_stdout_contains(
            "    read-manifest        Print a JSON representation of a Cargo.toml manifest.",
        )
        .run();
}

#[cargo_test]
fn list_builtin_aliases_with_descriptions() {
    let p = project().build();
    p.cargo("--list")
        .with_stdout_contains("    b                    alias: build")
        .with_stdout_contains("    c                    alias: check")
        .with_stdout_contains("    r                    alias: run")
        .with_stdout_contains("    t                    alias: test")
        .run();
}

#[cargo_test]
fn list_custom_aliases_with_descriptions() {
    let p = project_in_home("proj")
        .file(
            &paths::home().join(".cargo").join("config"),
            r#"
            [alias]
            myaliasstr = "foo --bar"
            myaliasvec = ["foo", "--bar"]
        "#,
        )
        .build();

    p.cargo("--list")
        .with_stdout_contains("    myaliasstr           alias: foo --bar")
        .with_stdout_contains("    myaliasvec           alias: foo --bar")
        .run();
}

#[cargo_test]
fn list_dedupe() {
    let p = project()
        .executable(Path::new("path-test-1").join("cargo-dupe"), "")
        .executable(Path::new("path-test-2").join("cargo-dupe"), "")
        .build();

    let mut path = path();
    path.push(p.root().join("path-test-1"));
    path.push(p.root().join("path-test-2"));
    let path = env::join_paths(path.iter()).unwrap();

    p.cargo("--list")
        .env("PATH", &path)
        .with_stdout_contains_n("    dupe", 1)
        .run();
}

#[cargo_test]
fn list_command_looks_at_path() {
    let proj = project()
        .executable(Path::new("path-test").join("cargo-1"), "")
        .build();

    let mut path = path();
    path.push(proj.root().join("path-test"));
    let path = env::join_paths(path.iter()).unwrap();
    let output = cargo_process("-v --list")
        .env("PATH", &path)
        .exec_with_output()
        .unwrap();
    let output = str::from_utf8(&output.stdout).unwrap();
    assert!(
        output.contains("\n    1                   "),
        "missing 1: {}",
        output
    );
}

#[cargo_test]
fn list_command_handles_known_external_commands() {
    let p = project()
        .executable(Path::new("path-test").join("cargo-fmt"), "")
        .build();

    let fmt_desc = "    fmt                  Formats all bin and lib files of the current crate using rustfmt.";

    // Without path - fmt isn't there
    p.cargo("--list")
        .env("PATH", "")
        .with_stdout_does_not_contain(fmt_desc)
        .run();

    // With path - fmt is there with known description
    let mut path = path();
    path.push(p.root().join("path-test"));
    let path = env::join_paths(path.iter()).unwrap();

    p.cargo("--list")
        .env("PATH", &path)
        .with_stdout_contains(fmt_desc)
        .run();
}

#[cargo_test]
fn list_command_resolves_symlinks() {
    let proj = project()
        .symlink(cargo_exe(), Path::new("path-test").join("cargo-2"))
        .build();

    let mut path = path();
    path.push(proj.root().join("path-test"));
    let path = env::join_paths(path.iter()).unwrap();
    let output = cargo_process("-v --list")
        .env("PATH", &path)
        .exec_with_output()
        .unwrap();
    let output = str::from_utf8(&output.stdout).unwrap();
    assert!(
        output.contains("\n    2                   "),
        "missing 2: {}",
        output
    );
}

#[cargo_test]
fn find_closest_capital_c_to_c() {
    cargo_process("C")
        .with_status(101)
        .with_stderr_contains(
            "\
error: no such subcommand: `C`

<tab>Did you mean `c`?
",
        )
        .run();
}

#[cargo_test]
fn find_closest_capital_b_to_b() {
    cargo_process("B")
        .with_status(101)
        .with_stderr_contains(
            "\
error: no such subcommand: `B`

<tab>Did you mean `b`?
",
        )
        .run();
}

#[cargo_test]
fn find_closest_biuld_to_build() {
    cargo_process("biuld")
        .with_status(101)
        .with_stderr_contains(
            "\
error: no such subcommand: `biuld`

<tab>Did you mean `build`?
",
        )
        .run();

    // But, if we actually have `biuld`, it must work!
    // https://github.com/rust-lang/cargo/issues/5201
    Package::new("cargo-biuld", "1.0.0")
        .file(
            "src/main.rs",
            r#"
                fn main() {
                    println!("Similar, but not identical to, build");
                }
            "#,
        )
        .publish();

    cargo_process("install cargo-biuld").run();
    cargo_process("biuld")
        .with_stdout("Similar, but not identical to, build\n")
        .run();
    cargo_process("--list")
        .with_stdout_contains(
            "    build                Compile a local package and all of its dependencies\n",
        )
        .with_stdout_contains("    biuld\n")
        .run();
}

#[cargo_test]
fn find_closest_alias() {
    let root = paths::root();
    let my_home = root.join("my_home");
    fs::create_dir(&my_home).unwrap();
    fs::write(
        &my_home.join("config"),
        r#"
            [alias]
            myalias = "build"
        "#,
    )
    .unwrap();

    cargo_process("myalais")
        .env("CARGO_HOME", &my_home)
        .with_status(101)
        .with_stderr_contains(
            "\
error: no such subcommand: `myalais`

<tab>Did you mean `myalias`?
",
        )
        .run();

    // But, if no alias is defined, it must not suggest one!
    cargo_process("myalais")
        .with_status(101)
        .with_stderr_contains(
            "\
error: no such subcommand: `myalais`
",
        )
        .with_stderr_does_not_contain(
            "\
<tab>Did you mean `myalias`?
",
        )
        .run();
}

// If a subcommand is more than an edit distance of 3 away, we don't make a suggestion.
#[cargo_test]
fn find_closest_dont_correct_nonsense() {
    cargo_process("there-is-no-way-that-there-is-a-command-close-to-this")
        .cwd(&paths::root())
        .with_status(101)
        .with_stderr(
            "\
[ERROR] no such subcommand: `there-is-no-way-that-there-is-a-command-close-to-this`

<tab>View all installed commands with `cargo --list`",
        )
        .run();
}

#[cargo_test]
fn displays_subcommand_on_error() {
    cargo_process("invalid-command")
        .with_status(101)
        .with_stderr(
            "\
[ERROR] no such subcommand: `invalid-command`

<tab>View all installed commands with `cargo --list`",
        )
        .run();
}

#[cargo_test]
fn override_cargo_home() {
    let root = paths::root();
    let my_home = root.join("my_home");
    fs::create_dir(&my_home).unwrap();
    fs::write(
        &my_home.join("config"),
        r#"
            [cargo-new]
            vcs = "none"
        "#,
    )
    .unwrap();

    cargo_process("new foo").env("CARGO_HOME", &my_home).run();

    assert!(!paths::root().join("foo/.git").is_dir());

    cargo_process("new foo2").run();

    assert!(paths::root().join("foo2/.git").is_dir());
}

#[cargo_test]
fn cargo_subcommand_env() {
    let src = format!(
        r#"
        use std::env;

        fn main() {{
            println!("{{}}", env::var("{}").unwrap());
        }}
        "#,
        cargo::CARGO_ENV
    );

    let p = project()
        .at("cargo-envtest")
        .file("Cargo.toml", &basic_bin_manifest("cargo-envtest"))
        .file("src/main.rs", &src)
        .build();

    let target_dir = p.target_debug_dir();

    p.cargo("build").run();
    assert!(p.bin("cargo-envtest").is_file());

    let cargo = cargo_exe().canonicalize().unwrap();
    let mut path = path();
    path.push(target_dir);
    let path = env::join_paths(path.iter()).unwrap();

    cargo_process("envtest")
        .env("PATH", &path)
        .with_stdout(cargo.to_str().unwrap())
        .run();
}

#[cargo_test]
fn cargo_subcommand_args() {
    let p = echo_subcommand();
    let cargo_foo_bin = p.bin("cargo-echo");
    assert!(cargo_foo_bin.is_file());

    let mut path = path();
    path.push(p.target_debug_dir());
    let path = env::join_paths(path.iter()).unwrap();

    cargo_process("echo bar -v --help")
        .env("PATH", &path)
        .with_stdout("echo bar -v --help")
        .run();
}

#[cargo_test]
fn explain() {
    cargo_process("--explain E0001")
        .with_stdout_contains(
            "This error suggests that the expression arm corresponding to the noted pattern",
        )
        .run();
}

#[cargo_test]
fn closed_output_ok() {
    // Checks that closed output doesn't cause an error.
    let mut p = cargo_process("--list").build_command();
    p.stdout(Stdio::piped()).stderr(Stdio::piped());
    let mut child = p.spawn().unwrap();
    // Close stdout
    drop(child.stdout.take());
    // Read stderr
    let mut s = String::new();
    child
        .stderr
        .as_mut()
        .unwrap()
        .read_to_string(&mut s)
        .unwrap();
    let status = child.wait().unwrap();
    assert!(status.success());
    assert!(s.is_empty(), "{}", s);
}
<|endoftext|>


<commit_before>
// ====--------------------------------------------------------------------------------------====//
//
// Parser for .cton files.
//
// ====--------------------------------------------------------------------------------------====//

use std::result;
use lexer::{self, Lexer, Token};
use cretonne::{types, repr};

pub use lexer::Location;

/// A parse error is returned when the parse failed.
pub struct Error {
    pub location: Location,
    pub message: String,
}

pub type Result<T> = result::Result<T, Error>;

pub struct Parser<'a> {
    lex: Lexer<'a>,

    lex_error: Option<lexer::Error>,

    // Current lookahead token.
    lookahead: Option<Token<'a>>,

    // Location of lookahead.
    location: Location,
}

impl<'a> Parser<'a> {
    // Consume the current lookahead token and return it.
    fn consume(&mut self) -> Token<'a> {
        self.lookahead.take().expect("No token to consume")
    }

    // Get the current lookahead token, after making sure there is one.
    fn token(&mut self) -> Option<Token<'a>> {
        if self.lookahead == None {
            match self.lex.next() {
                Some(Ok(lexer::LocatedToken { token, location })) => {
                    self.lookahead = Some(token);
                    self.location = location;
                }
                Some(Err(lexer::LocatedError { error, location })) => {
                    self.lex_error = Some(error);
                    self.location = location;
                }
                None => {}
            }
        }
        return self.lookahead;
    }

    // Generate an error.
    fn error(&self, message: &str) -> Error {
        Error {
            location: self.location,
            message:
                // If we have a lexer error latched, report that.
                match self.lex_error {
                    Some(lexer::Error::InvalidChar) => "Invalid character".to_string(),
                    None => message.to_string(),
                }
        }
    }

    // Match and consume a token without payload.
    fn match_token(&mut self, want: Token<'a>, err_msg: &str) -> Result<Token<'a>> {
        if self.token() == Some(want) {
            Ok(self.consume())
        } else {
            Err(self.error(err_msg))
        }
    }

    // If the next token is a `want`, consume it, otherwise do nothing.
    fn optional(&mut self, want: Token<'a>) -> bool {
        if self.token() == Some(want) {
            self.consume();
            true
        } else {
            false
        }
    }

    // Parse a whole function definition.
    //
    // function ::= * "function" name signature { ... }
    //
    fn parse_function(&mut self) -> Result<repr::Function> {
        try!(self.match_token(Token::Function, "Expected 'function' keyword"));

        // function ::= "function" * name signature { ... }
        let name = try!(self.parse_function_name());

        // function ::= "function" name * signature { ... }
        let sig = try!(self.parse_signature());

        let mut func = repr::Function::new();

        try!(self.match_token(Token::LBrace, "Expected '{' before function body"));
        try!(self.match_token(Token::RBrace, "Expected '}' after function body"));

        Ok(func)
    }

    // Parse a function name.
    //
    // function ::= "function" * name signature { ... }
    //
    fn parse_function_name(&mut self) -> Result<String> {
        match self.token() {
            Some(Token::Identifier(s)) => {
                self.consume();
                Ok(s.to_string())
            }
            _ => Err(self.error("Expected function name")),
        }
    }

    // Parse a function signature.
    //
    // signature ::=  * "(" [arglist] ")" ["->" retlist] [call_conv]
    //
    fn parse_signature(&mut self) -> Result<types::Signature> {
        let mut sig = types::Signature::new();

        try!(self.match_token(Token::LPar, "Expected function signature: '(' args... ')'"));
        // signature ::=  "(" * [arglist] ")" ["->" retlist] [call_conv]
        if self.token() != Some(Token::RPar) {
            sig.argument_types = try!(self.parse_argument_list());
        }
        try!(self.match_token(Token::RPar, "Expected ')' after function arguments"));
        if self.optional(Token::Arrow) {
            sig.return_types = try!(self.parse_argument_list());
        }

        // TBD: calling convention.

        Ok(sig)
    }

    // Parse list of function argument / return value types.
    //
    // arglist ::= * arg { "," arg }
    //
    fn parse_argument_list(&mut self) -> Result<Vec<types::ArgumentType>> {
        let mut list = Vec::new();

        // arglist ::= * arg { "," arg }
        list.push(try!(self.parse_argument_type()));

        // arglist ::= arg * { "," arg }
        while self.optional(Token::Comma) {
            // arglist ::= arg { "," * arg }
            list.push(try!(self.parse_argument_type()));
        }

        Ok(list)
    }

    // Parse a single argument type with flags.
    fn parse_argument_type(&mut self) -> Result<types::ArgumentType> {
        // arg ::= * type { flag }
        let mut arg = if let Some(Token::Type(t)) = self.token() {
            types::ArgumentType::new(t)
        } else {
            return Err(self.error("Expected argument type"));
        };
        self.consume();

        // arg ::= type * { flag }
        while let Some(Token::Identifier(s)) = self.token() {
            match s {
                "uext" => arg.extension = types::ArgumentExtension::Uext,
                "sext" => arg.extension = types::ArgumentExtension::Sext,
                "inreg" => arg.inreg = true,
                _ => break,
            }
            self.consume();
        }

        Ok(arg)
    }
}
<commit_msg>Begin parser unit tests, add public interface.<commit_after>
// ====--------------------------------------------------------------------------------------====//
//
// Parser for .cton files.
//
// ====--------------------------------------------------------------------------------------====//

use std::result;
use lexer::{self, Lexer, Token};
use cretonne::{types, repr};

pub use lexer::Location;

/// A parse error is returned when the parse failed.
#[derive(Debug)]
pub struct Error {
    pub location: Location,
    pub message: String,
}

pub type Result<T> = result::Result<T, Error>;

pub struct Parser<'a> {
    lex: Lexer<'a>,

    lex_error: Option<lexer::Error>,

    // Current lookahead token.
    lookahead: Option<Token<'a>>,

    // Location of lookahead.
    location: Location,
}

impl<'a> Parser<'a> {
    /// Create a new `Parser` which reads `text`. The referenced text must outlive the parser.
    pub fn new(text: &'a str) -> Parser {
        Parser {
            lex: Lexer::new(text),
            lex_error: None,
            lookahead: None,
            location: Location { line_number: 0 },
        }
    }

    /// Parse the entire string into a list of functions.
    pub fn parse(text: &'a str) -> Result<Vec<repr::Function>> {
        Self::new(text).parse_function_list()
    }

    // Consume the current lookahead token and return it.
    fn consume(&mut self) -> Token<'a> {
        self.lookahead.take().expect("No token to consume")
    }

    // Get the current lookahead token, after making sure there is one.
    fn token(&mut self) -> Option<Token<'a>> {
        if self.lookahead == None {
            match self.lex.next() {
                Some(Ok(lexer::LocatedToken { token, location })) => {
                    self.lookahead = Some(token);
                    self.location = location;
                }
                Some(Err(lexer::LocatedError { error, location })) => {
                    self.lex_error = Some(error);
                    self.location = location;
                }
                None => {}
            }
        }
        return self.lookahead;
    }

    // Generate an error.
    fn error(&self, message: &str) -> Error {
        Error {
            location: self.location,
            message:
                // If we have a lexer error latched, report that.
                match self.lex_error {
                    Some(lexer::Error::InvalidChar) => "invalid character".to_string(),
                    None => message.to_string(),
                }
        }
    }

    // Match and consume a token without payload.
    fn match_token(&mut self, want: Token<'a>, err_msg: &str) -> Result<Token<'a>> {
        if self.token() == Some(want) {
            Ok(self.consume())
        } else {
            Err(self.error(err_msg))
        }
    }

    // If the next token is a `want`, consume it, otherwise do nothing.
    fn optional(&mut self, want: Token<'a>) -> bool {
        if self.token() == Some(want) {
            self.consume();
            true
        } else {
            false
        }
    }

    /// Parse a list of function definitions.
    ///
    /// This is the top-level parse function matching the whole contents of a file.
    pub fn parse_function_list(&mut self) -> Result<Vec<repr::Function>> {
        let mut list = Vec::new();
        while self.token().is_some() {
            list.push(try!(self.parse_function()));
        }
        Ok(list)
    }

    // Parse a whole function definition.
    //
    // function ::= * "function" name signature { ... }
    //
    fn parse_function(&mut self) -> Result<repr::Function> {
        try!(self.match_token(Token::Function, "expected 'function' keyword"));

        // function ::= "function" * name signature { ... }
        let name = try!(self.parse_function_name());

        // function ::= "function" name * signature { ... }
        let sig = try!(self.parse_signature());

        let mut func = repr::Function::new();

        try!(self.match_token(Token::LBrace, "expected '{' before function body"));
        try!(self.match_token(Token::RBrace, "expected '}' after function body"));

        Ok(func)
    }

    // Parse a function name.
    //
    // function ::= "function" * name signature { ... }
    //
    fn parse_function_name(&mut self) -> Result<String> {
        match self.token() {
            Some(Token::Identifier(s)) => {
                self.consume();
                Ok(s.to_string())
            }
            _ => Err(self.error("expected function name")),
        }
    }

    // Parse a function signature.
    //
    // signature ::=  * "(" [arglist] ")" ["->" retlist] [call_conv]
    //
    fn parse_signature(&mut self) -> Result<types::Signature> {
        let mut sig = types::Signature::new();

        try!(self.match_token(Token::LPar, "expected function signature: '(' args... ')'"));
        // signature ::=  "(" * [arglist] ")" ["->" retlist] [call_conv]
        if self.token() != Some(Token::RPar) {
            sig.argument_types = try!(self.parse_argument_list());
        }
        try!(self.match_token(Token::RPar, "expected ')' after function arguments"));
        if self.optional(Token::Arrow) {
            sig.return_types = try!(self.parse_argument_list());
        }

        // TBD: calling convention.

        Ok(sig)
    }

    // Parse list of function argument / return value types.
    //
    // arglist ::= * arg { "," arg }
    //
    fn parse_argument_list(&mut self) -> Result<Vec<types::ArgumentType>> {
        let mut list = Vec::new();

        // arglist ::= * arg { "," arg }
        list.push(try!(self.parse_argument_type()));

        // arglist ::= arg * { "," arg }
        while self.optional(Token::Comma) {
            // arglist ::= arg { "," * arg }
            list.push(try!(self.parse_argument_type()));
        }

        Ok(list)
    }

    // Parse a single argument type with flags.
    fn parse_argument_type(&mut self) -> Result<types::ArgumentType> {
        // arg ::= * type { flag }
        let mut arg = if let Some(Token::Type(t)) = self.token() {
            types::ArgumentType::new(t)
        } else {
            return Err(self.error("expected argument type"));
        };
        self.consume();

        // arg ::= type * { flag }
        while let Some(Token::Identifier(s)) = self.token() {
            match s {
                "uext" => arg.extension = types::ArgumentExtension::Uext,
                "sext" => arg.extension = types::ArgumentExtension::Sext,
                "inreg" => arg.inreg = true,
                _ => break,
            }
            self.consume();
        }

        Ok(arg)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use cretonne::types::{self, ArgumentType, ArgumentExtension};

    #[test]
    fn argument_type() {
        let mut p = Parser::new("i32 sext");
        let arg = p.parse_argument_type().unwrap();
        assert_eq!(arg,
                   ArgumentType {
                       value_type: types::I32,
                       extension: ArgumentExtension::Sext,
                       inreg: false,
                   });
        let Error { location, message } = p.parse_argument_type().unwrap_err();
        assert_eq!(location.line_number, 1);
        assert_eq!(message, "expected argument type");
    }
}
<|endoftext|>


<commit_before>// Copyright 2015 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.

#![unstable(reason = "not public", issue = "0", feature = "fd")]

use io::{self, Read};
use libc::{self, c_int, size_t, c_void};
use mem;
use sync::atomic::{AtomicBool, Ordering};
use sys::cvt;
use sys_common::AsInner;
use sys_common::io::read_to_end_uninitialized;

pub struct FileDesc {
    fd: c_int,
}

impl FileDesc {
    pub fn new(fd: c_int) -> FileDesc {
        FileDesc { fd: fd }
    }

    pub fn raw(&self) -> c_int { self.fd }

    /// Extracts the actual filedescriptor without closing it.
    pub fn into_raw(self) -> c_int {
        let fd = self.fd;
        mem::forget(self);
        fd
    }

    pub fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
        let ret = cvt(unsafe {
            libc::read(self.fd,
                       buf.as_mut_ptr() as *mut c_void,
                       buf.len() as size_t)
        })?;
        Ok(ret as usize)
    }

    pub fn read_to_end(&self, buf: &mut Vec<u8>) -> io::Result<usize> {
        let mut me = self;
        (&mut me).read_to_end(buf)
    }

    pub fn write(&self, buf: &[u8]) -> io::Result<usize> {
        let ret = cvt(unsafe {
            libc::write(self.fd,
                        buf.as_ptr() as *const c_void,
                        buf.len() as size_t)
        })?;
        Ok(ret as usize)
    }

    #[cfg(not(any(target_env = "newlib", target_os = "solaris", target_os = "emscripten")))]
    pub fn set_cloexec(&self) -> io::Result<()> {
        unsafe {
            cvt(libc::ioctl(self.fd, libc::FIOCLEX))?;
            Ok(())
        }
    }
    #[cfg(any(target_env = "newlib", target_os = "solaris", target_os = "emscripten"))]
    pub fn set_cloexec(&self) -> io::Result<()> {
        unsafe {
            let previous = cvt(libc::fcntl(self.fd, libc::F_GETFD))?;
            cvt(libc::fcntl(self.fd, libc::F_SETFD, previous | libc::FD_CLOEXEC))?;
            Ok(())
        }
    }

    pub fn set_nonblocking(&self, nonblocking: bool) -> io::Result<()> {
        unsafe {
            let previous = cvt(libc::fcntl(self.fd, libc::F_GETFL))?;
            let new = if nonblocking {
                previous | libc::O_NONBLOCK
            } else {
                previous & !libc::O_NONBLOCK
            };
            cvt(libc::fcntl(self.fd, libc::F_SETFL, new))?;
            Ok(())
        }
    }

    pub fn duplicate(&self) -> io::Result<FileDesc> {
        // We want to atomically duplicate this file descriptor and set the
        // CLOEXEC flag, and currently that's done via F_DUPFD_CLOEXEC. This
        // flag, however, isn't supported on older Linux kernels (earlier than
        // 2.6.24).
        //
        // To detect this and ensure that CLOEXEC is still set, we
        // follow a strategy similar to musl [1] where if passing
        // F_DUPFD_CLOEXEC causes `fcntl` to return EINVAL it means it's not
        // supported (the third parameter, 0, is always valid), so we stop
        // trying that.
        //
        // Also note that Android doesn't have F_DUPFD_CLOEXEC, but get it to
        // resolve so we at least compile this.
        //
        // [1]: http://comments.gmane.org/gmane.linux.lib.musl.general/2963
        #[cfg(any(target_os = "android", target_os = "haiku"))]
        use libc::F_DUPFD as F_DUPFD_CLOEXEC;
        #[cfg(not(any(target_os = "android", target_os="haiku")))]
        use libc::F_DUPFD_CLOEXEC;

        let make_filedesc = |fd| {
            let fd = FileDesc::new(fd);
            fd.set_cloexec()?;
            Ok(fd)
        };
        static TRY_CLOEXEC: AtomicBool =
            AtomicBool::new(!cfg!(target_os = "android"));
        let fd = self.raw();
        if TRY_CLOEXEC.load(Ordering::Relaxed) {
            match cvt(unsafe { libc::fcntl(fd, F_DUPFD_CLOEXEC, 0) }) {
                // We *still* call the `set_cloexec` method as apparently some
                // linux kernel at some point stopped setting CLOEXEC even
                // though it reported doing so on F_DUPFD_CLOEXEC.
                Ok(fd) => {
                    return Ok(if cfg!(target_os = "linux") {
                        make_filedesc(fd)?
                    } else {
                        FileDesc::new(fd)
                    })
                }
                Err(ref e) if e.raw_os_error() == Some(libc::EINVAL) => {
                    TRY_CLOEXEC.store(false, Ordering::Relaxed);
                }
                Err(e) => return Err(e),
            }
        }
        cvt(unsafe { libc::fcntl(fd, libc::F_DUPFD, 0) }).and_then(make_filedesc)
    }
}

impl<'a> Read for &'a FileDesc {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        (**self).read(buf)
    }

    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
        unsafe { read_to_end_uninitialized(self, buf) }
    }
}

impl AsInner<c_int> for FileDesc {
    fn as_inner(&self) -> &c_int { &self.fd }
}

impl Drop for FileDesc {
    fn drop(&mut self) {
        // Note that errors are ignored when closing a file descriptor. The
        // reason for this is that if an error occurs we don't actually know if
        // the file descriptor was closed or not, and if we retried (for
        // something like EINTR), we might close another valid file descriptor
        // (opened after we closed ours.
        let _ = unsafe { libc::close(self.fd) };
    }
}
<commit_msg>Haiku: Work around the lack of the FIOCLEX ioctl<commit_after>// Copyright 2015 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.

#![unstable(reason = "not public", issue = "0", feature = "fd")]

use io::{self, Read};
use libc::{self, c_int, size_t, c_void};
use mem;
use sync::atomic::{AtomicBool, Ordering};
use sys::cvt;
use sys_common::AsInner;
use sys_common::io::read_to_end_uninitialized;

pub struct FileDesc {
    fd: c_int,
}

impl FileDesc {
    pub fn new(fd: c_int) -> FileDesc {
        FileDesc { fd: fd }
    }

    pub fn raw(&self) -> c_int { self.fd }

    /// Extracts the actual filedescriptor without closing it.
    pub fn into_raw(self) -> c_int {
        let fd = self.fd;
        mem::forget(self);
        fd
    }

    pub fn read(&self, buf: &mut [u8]) -> io::Result<usize> {
        let ret = cvt(unsafe {
            libc::read(self.fd,
                       buf.as_mut_ptr() as *mut c_void,
                       buf.len() as size_t)
        })?;
        Ok(ret as usize)
    }

    pub fn read_to_end(&self, buf: &mut Vec<u8>) -> io::Result<usize> {
        let mut me = self;
        (&mut me).read_to_end(buf)
    }

    pub fn write(&self, buf: &[u8]) -> io::Result<usize> {
        let ret = cvt(unsafe {
            libc::write(self.fd,
                        buf.as_ptr() as *const c_void,
                        buf.len() as size_t)
        })?;
        Ok(ret as usize)
    }

    #[cfg(not(any(target_env = "newlib",
                  target_os = "solaris",
                  target_os = "emscripten",
                  target_os = "haiku")))]
    pub fn set_cloexec(&self) -> io::Result<()> {
        unsafe {
            cvt(libc::ioctl(self.fd, libc::FIOCLEX))?;
            Ok(())
        }
    }
    #[cfg(any(target_env = "newlib",
              target_os = "solaris",
              target_os = "emscripten",
              target_os = "haiku"))]
    pub fn set_cloexec(&self) -> io::Result<()> {
        unsafe {
            let previous = cvt(libc::fcntl(self.fd, libc::F_GETFD))?;
            cvt(libc::fcntl(self.fd, libc::F_SETFD, previous | libc::FD_CLOEXEC))?;
            Ok(())
        }
    }

    pub fn set_nonblocking(&self, nonblocking: bool) -> io::Result<()> {
        unsafe {
            let previous = cvt(libc::fcntl(self.fd, libc::F_GETFL))?;
            let new = if nonblocking {
                previous | libc::O_NONBLOCK
            } else {
                previous & !libc::O_NONBLOCK
            };
            cvt(libc::fcntl(self.fd, libc::F_SETFL, new))?;
            Ok(())
        }
    }

    pub fn duplicate(&self) -> io::Result<FileDesc> {
        // We want to atomically duplicate this file descriptor and set the
        // CLOEXEC flag, and currently that's done via F_DUPFD_CLOEXEC. This
        // flag, however, isn't supported on older Linux kernels (earlier than
        // 2.6.24).
        //
        // To detect this and ensure that CLOEXEC is still set, we
        // follow a strategy similar to musl [1] where if passing
        // F_DUPFD_CLOEXEC causes `fcntl` to return EINVAL it means it's not
        // supported (the third parameter, 0, is always valid), so we stop
        // trying that.
        //
        // Also note that Android doesn't have F_DUPFD_CLOEXEC, but get it to
        // resolve so we at least compile this.
        //
        // [1]: http://comments.gmane.org/gmane.linux.lib.musl.general/2963
        #[cfg(any(target_os = "android", target_os = "haiku"))]
        use libc::F_DUPFD as F_DUPFD_CLOEXEC;
        #[cfg(not(any(target_os = "android", target_os="haiku")))]
        use libc::F_DUPFD_CLOEXEC;

        let make_filedesc = |fd| {
            let fd = FileDesc::new(fd);
            fd.set_cloexec()?;
            Ok(fd)
        };
        static TRY_CLOEXEC: AtomicBool =
            AtomicBool::new(!cfg!(target_os = "android"));
        let fd = self.raw();
        if TRY_CLOEXEC.load(Ordering::Relaxed) {
            match cvt(unsafe { libc::fcntl(fd, F_DUPFD_CLOEXEC, 0) }) {
                // We *still* call the `set_cloexec` method as apparently some
                // linux kernel at some point stopped setting CLOEXEC even
                // though it reported doing so on F_DUPFD_CLOEXEC.
                Ok(fd) => {
                    return Ok(if cfg!(target_os = "linux") {
                        make_filedesc(fd)?
                    } else {
                        FileDesc::new(fd)
                    })
                }
                Err(ref e) if e.raw_os_error() == Some(libc::EINVAL) => {
                    TRY_CLOEXEC.store(false, Ordering::Relaxed);
                }
                Err(e) => return Err(e),
            }
        }
        cvt(unsafe { libc::fcntl(fd, libc::F_DUPFD, 0) }).and_then(make_filedesc)
    }
}

impl<'a> Read for &'a FileDesc {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        (**self).read(buf)
    }

    fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
        unsafe { read_to_end_uninitialized(self, buf) }
    }
}

impl AsInner<c_int> for FileDesc {
    fn as_inner(&self) -> &c_int { &self.fd }
}

impl Drop for FileDesc {
    fn drop(&mut self) {
        // Note that errors are ignored when closing a file descriptor. The
        // reason for this is that if an error occurs we don't actually know if
        // the file descriptor was closed or not, and if we retried (for
        // something like EINTR), we might close another valid file descriptor
        // (opened after we closed ours.
        let _ = unsafe { libc::close(self.fd) };
    }
}
<|endoftext|>