<commit_before>use tf;
use libc::{c_char, c_int};
use std::ffi::CString;
use std::marker;
use std::path::Path;
use std::ptr;
use super::Code;
use super::DataType;
use super::Graph;
use super::GraphTrait;
use super::Operation;
use super::OperationTrait;
use super::Result;
use super::SessionOptions;
use super::Status;
use super::Tensor;
use super::TensorType;

/// Manages a single graph and execution.
#[derive(Debug)]
pub struct Session {
    inner: *mut tf::TF_Session,
}

impl Session {
    /// Creates a session.
    pub fn new(options: &SessionOptions, graph: &Graph) -> Result<Self> {
        let mut status = Status::new();
        let inner = unsafe { tf::TF_NewSession(graph.inner(), options.inner, status.inner()) };
        if inner.is_null() {
            Err(status)
        } else {
            Ok(Session { inner: inner })
        }
    }

    /// Loads a session from an exported model.
    pub fn from_saved_model<P: AsRef<Path>, Tag: AsRef<str>, Tags: IntoIterator<Item = Tag>>
        (options: &SessionOptions,
         tags: Tags,
         graph: &mut Graph,
         export_dir: P)
         -> Result<Self> {
        let mut status = Status::new();

        let export_dir_cstr =
            try!(export_dir.as_ref()
                     .to_str()
                     .and_then(|s| CString::new(s.as_bytes()).ok())
                     .ok_or_else(|| invalid_arg!("Invalid export directory path")));

        let tags_cstr: Vec<_> = try!(tags.into_iter()
                                         .map(|t| CString::new(t.as_ref()))
                                         .collect::<::std::result::Result<_, _>>()
                                         .map_err(|_| invalid_arg!("Invalid tag name")));
        // keeping tags_cstr to retain strings in memory
        let tags_ptr: Vec<*const c_char> = tags_cstr.iter().map(|t| t.as_ptr()).collect();

        let inner = unsafe {
            tf::TF_LoadSessionFromSavedModel(options.inner,
                                             ptr::null(),
                                             export_dir_cstr.as_ptr(),
                                             tags_ptr.as_ptr(),
                                             tags_ptr.len() as c_int,
                                             graph.inner(),
                                             ptr::null_mut(),
                                             status.inner())
        };
        if inner.is_null() {
            Err(status)
        } else {
            Ok(Session { inner: inner })
        }
    }

    /// Closes the session.
    pub fn close(&mut self) -> Result<()> {
        let mut status = Status::new();
        unsafe {
            tf::TF_CloseSession(self.inner, status.inner());
        }
        status.into_result()
    }

    /// Runs the graph, feeding the inputs and then fetching the outputs requested in the step.
    pub fn run(&mut self, step: &mut StepWithGraph) -> Result<()> {
        // Copy the input tensors because TF_Run consumes them.
        let mut input_tensors = Vec::with_capacity(step.input_tensors.len());
        for &input_tensor in &step.input_tensors {
            unsafe {
                let mut dims = Vec::with_capacity(tf::TF_NumDims(input_tensor) as usize);
                for i in 0..dims.capacity() {
                    dims.push(tf::TF_Dim(input_tensor, i as c_int));
                }
                input_tensors.push(tf::TF_NewTensor(tf::TF_TensorType(input_tensor),
                                                    dims.as_ptr(),
                                                    dims.len() as c_int,
                                                    tf::TF_TensorData(input_tensor),
                                                    tf::TF_TensorByteSize(input_tensor),
                                                    Some(super::noop_deallocator),
                                                    ptr::null_mut()));
            }
        }

        // In case we're running it a second time and not all outputs were taken out.
        step.drop_output_tensors();

        let mut status = Status::new();
        unsafe {
            tf::TF_SessionRun(self.inner,
                              ptr::null(),
                              step.input_ports.as_ptr(),
                              input_tensors.as_ptr() as *const *const tf::TF_Tensor,
                              input_tensors.len() as c_int,
                              step.output_ports.as_ptr(),
                              step.output_tensors.as_mut_ptr(),
                              step.output_tensors.len() as c_int,
                              step.target_operations.as_mut_ptr(),
                              step.target_operations.len() as c_int,
                              ptr::null_mut(),
                              status.inner());
        };
        status.into_result()
    }
}

impl Drop for Session {
    fn drop(&mut self) {
        let mut status = Status::new();
        unsafe {
            tf::TF_DeleteSession(self.inner, status.inner());
        }
        // TODO: What do we do with the status?
    }
}

////////////////////////

/// An opaque token for retrieving an output from a computation.
#[derive(Copy,Clone,Debug)]
pub struct OutputToken {
    index: usize,
}

/// Manages the inputs and outputs for a single execution of a graph.
///
/// Typical usage involves creating an instance of this struct,
/// adding some inputs to it, requesting some outputs, passing it to `Session::run`
/// and then taking the outputs out of it.
///
/// This will be renamed to Step once the old API goes away.
#[derive(Debug)]
pub struct StepWithGraph<'l> {
    input_ports: Vec<tf::TF_Output>,
    input_tensors: Vec<*mut tf::TF_Tensor>,

    output_ports: Vec<tf::TF_Output>,
    output_tensors: Vec<*mut tf::TF_Tensor>,

    target_operations: Vec<*const tf::TF_Operation>,

    phantom: marker::PhantomData<&'l ()>,
}

impl<'l> StepWithGraph<'l> {
    /// Creates a StepWithGraph.
    pub fn new() -> Self {
        StepWithGraph {
            input_ports: vec![],
            input_tensors: vec![],

            output_ports: vec![],
            output_tensors: vec![],

            target_operations: vec![],

            phantom: marker::PhantomData,
        }
    }

    /// Adds an input to be fed to the graph.
    pub fn add_input<T: TensorType>(&mut self,
                                    operation: &Operation,
                                    index: c_int,
                                    tensor: &'l Tensor<T>) {
        self.input_ports.push(tf::TF_Output {
                                  oper: operation.inner(),
                                  index: index,
                              });
        self.input_tensors.push(tensor.inner);
    }

    /// Requests that an output is fetched from the graph after running this step.
    /// Returns an index that you can then use to fetch this output from the step after running it.
    pub fn request_output(&mut self, operation: &Operation, index: c_int) -> OutputToken {
        self.output_ports.push(tf::TF_Output {
                                   oper: operation.inner(),
                                   index: index,
                               });
        self.output_tensors.push(ptr::null_mut());
        OutputToken { index: self.output_tensors.len() - 1 }
    }

    /// Extracts a tensor output given an index. A given index can only be extracted
    /// once per `Session::run`.
    /// Returns an error if output_idx is out of range, output is unavailable or the
    /// requested type does not match the type of the actual tensor.
    pub fn take_output<T: TensorType>(&mut self, token: OutputToken) -> Result<Tensor<T>> {
        let output_idx = token.index;
        if output_idx >= self.output_tensors.len() {
            return Err(Status::new_set(Code::OutOfRange,
                                       &format!("Requested output index is out of range: {} vs \
                                                 {}",
                                                output_idx,
                                                self.output_tensors.len()))
                               .unwrap());
        }
        if self.output_tensors[output_idx].is_null() {
            return Err(Status::new_set(Code::Unavailable,
                                       "Output not available. Either it was already taken, or \
                                        this step has not been sucessfully run yet.")
                               .unwrap());
        }
        let actual_data_type = self.output_data_type(output_idx).unwrap();
        if actual_data_type != T::data_type() {
            return Err(invalid_arg!("Requested tensor type does not match actual tensor type: \
                                     {} vs {}",
                                    actual_data_type,
                                    T::data_type()));
        }
        let tensor = unsafe { Tensor::from_tf_tensor(self.output_tensors[output_idx]).unwrap() };
        self.output_tensors[output_idx] = ptr::null_mut();
        Ok(tensor)
    }

    /// Adds a target operation to be executed when running the graph.
    pub fn add_target(&mut self, operation: &Operation) {
        self.target_operations.push(operation.inner());
    }

    /// Retuns the type of the tensor given an index.
    /// Returns `None` if the index is out of range or the output is not yet available.
    pub fn output_data_type(&self, output_idx: usize) -> Option<DataType> {
        if output_idx >= self.output_tensors.len() {
            return None;
        }
        if self.output_tensors[output_idx].is_null() {
            return None;
        }
        unsafe { Some(DataType::from_c(tf::TF_TensorType(self.output_tensors[output_idx]))) }
    }

    fn drop_output_tensors(&mut self) {
        for mut tensor in &mut self.output_tensors {
            // TODO: Is TF_DeleteTensor NULL safe?
            if !tensor.is_null() {
                unsafe {
                    tf::TF_DeleteTensor(*tensor);
                }
            }
            *tensor = ptr::null_mut();
        }
    }
}

impl<'l> Drop for StepWithGraph<'l> {
    fn drop(&mut self) {
        self.drop_output_tensors();
    }
}

////////////////////////

#[cfg(test)]
mod tests {
    use super::*;
    use super::super::DataType;
    use super::super::Graph;
    use super::super::Operation;
    use super::super::Output;
    use super::super::SessionOptions;
    use super::super::Shape;
    use super::super::Tensor;

    fn create_session() -> (Session, Operation, Operation) {
        let mut g = Graph::new();
        let two = {
            let mut nd = g.new_operation("Const", "two").unwrap();
            nd.set_attr_type("dtype", DataType::Float).unwrap();
            let mut value = Tensor::new(&[1]);
            value[0] = 2.0f32;
            nd.set_attr_tensor("value", value).unwrap();
            nd.finish().unwrap()
        };
        let x = {
            let mut nd = g.new_operation("Placeholder", "x").unwrap();
            nd.set_attr_type("dtype", DataType::Float).unwrap();
            nd.set_attr_shape("shape", &Shape(Some(vec![]))).unwrap();
            nd.finish().unwrap()
        };
        let y = {
            let mut nd = g.new_operation("Mul", "y").unwrap();
            nd.add_input(Output {
                             operation: two,
                             index: 0,
                         });
            nd.add_input(Output {
                             operation: x.clone(),
                             index: 0,
                         });
            nd.finish().unwrap()
        };
        let options = SessionOptions::new();
        match Session::new(&options, &g) {
            Ok(session) => (session, x, y),
            Err(status) => panic!("Creating session failed with status: {}", status),
        }
    }

    #[test]
    fn smoke() {
        create_session();
    }

    #[test]
    fn test_close() {
        let (mut session, _, _) = create_session();
        let status = session.close();
        assert!(status.is_ok());
    }

    #[test]
    fn test_run() {
        let (mut session, x_operation, y_operation) = create_session();
        let mut x = <Tensor<f32>>::new(&[2]);
        x[0] = 2.0;
        x[1] = 3.0;
        let mut step = StepWithGraph::new();
        step.add_input(&x_operation, 0, &x);
        let output_token = step.request_output(&y_operation, 0);
        session.run(&mut step).unwrap();
        let output_tensor = step.take_output::<f32>(output_token).unwrap();
        assert_eq!(output_tensor.len(), 2);
        assert_eq!(output_tensor[0], 4.0);
        assert_eq!(output_tensor[1], 6.0);
    }
}
<commit_msg>TF_Run caller retains object ownership since 0.12.<commit_after>use tf;
use libc::{c_char, c_int};
use std::ffi::CString;
use std::marker;
use std::path::Path;
use std::ptr;
use super::Code;
use super::DataType;
use super::Graph;
use super::GraphTrait;
use super::Operation;
use super::OperationTrait;
use super::Result;
use super::SessionOptions;
use super::Status;
use super::Tensor;
use super::TensorType;

/// Manages a single graph and execution.
#[derive(Debug)]
pub struct Session {
    inner: *mut tf::TF_Session,
}

impl Session {
    /// Creates a session.
    pub fn new(options: &SessionOptions, graph: &Graph) -> Result<Self> {
        let mut status = Status::new();
        let inner = unsafe { tf::TF_NewSession(graph.inner(), options.inner, status.inner()) };
        if inner.is_null() {
            Err(status)
        } else {
            Ok(Session { inner: inner })
        }
    }

    /// Loads a session from an exported model.
    pub fn from_saved_model<P: AsRef<Path>, Tag: AsRef<str>, Tags: IntoIterator<Item = Tag>>
        (options: &SessionOptions,
         tags: Tags,
         graph: &mut Graph,
         export_dir: P)
         -> Result<Self> {
        let mut status = Status::new();

        let export_dir_cstr =
            try!(export_dir.as_ref()
                     .to_str()
                     .and_then(|s| CString::new(s.as_bytes()).ok())
                     .ok_or_else(|| invalid_arg!("Invalid export directory path")));

        let tags_cstr: Vec<_> = try!(tags.into_iter()
                                         .map(|t| CString::new(t.as_ref()))
                                         .collect::<::std::result::Result<_, _>>()
                                         .map_err(|_| invalid_arg!("Invalid tag name")));
        // keeping tags_cstr to retain strings in memory
        let tags_ptr: Vec<*const c_char> = tags_cstr.iter().map(|t| t.as_ptr()).collect();

        let inner = unsafe {
            tf::TF_LoadSessionFromSavedModel(options.inner,
                                             ptr::null(),
                                             export_dir_cstr.as_ptr(),
                                             tags_ptr.as_ptr(),
                                             tags_ptr.len() as c_int,
                                             graph.inner(),
                                             ptr::null_mut(),
                                             status.inner())
        };
        if inner.is_null() {
            Err(status)
        } else {
            Ok(Session { inner: inner })
        }
    }

    /// Closes the session.
    pub fn close(&mut self) -> Result<()> {
        let mut status = Status::new();
        unsafe {
            tf::TF_CloseSession(self.inner, status.inner());
        }
        status.into_result()
    }

    /// Runs the graph, feeding the inputs and then fetching the outputs requested in the step.
    pub fn run(&mut self, step: &mut StepWithGraph) -> Result<()> {
        // In case we're running it a second time and not all outputs were taken out.
        step.drop_output_tensors();

        let mut status = Status::new();
        unsafe {
            tf::TF_SessionRun(self.inner,
                              ptr::null(),
                              step.input_ports.as_ptr(),
                              step.input_tensors.as_ptr() as *const *const tf::TF_Tensor,
                              step.input_tensors.len() as c_int,
                              step.output_ports.as_ptr(),
                              step.output_tensors.as_mut_ptr(),
                              step.output_tensors.len() as c_int,
                              step.target_operations.as_mut_ptr(),
                              step.target_operations.len() as c_int,
                              ptr::null_mut(),
                              status.inner());
        };
        status.into_result()
    }
}

impl Drop for Session {
    fn drop(&mut self) {
        let mut status = Status::new();
        unsafe {
            tf::TF_DeleteSession(self.inner, status.inner());
        }
        // TODO: What do we do with the status?
    }
}

////////////////////////

/// An opaque token for retrieving an output from a computation.
#[derive(Copy,Clone,Debug)]
pub struct OutputToken {
    index: usize,
}

/// Manages the inputs and outputs for a single execution of a graph.
///
/// Typical usage involves creating an instance of this struct,
/// adding some inputs to it, requesting some outputs, passing it to `Session::run`
/// and then taking the outputs out of it.
///
/// This will be renamed to Step once the old API goes away.
#[derive(Debug)]
pub struct StepWithGraph<'l> {
    input_ports: Vec<tf::TF_Output>,
    input_tensors: Vec<*mut tf::TF_Tensor>,

    output_ports: Vec<tf::TF_Output>,
    output_tensors: Vec<*mut tf::TF_Tensor>,

    target_operations: Vec<*const tf::TF_Operation>,

    phantom: marker::PhantomData<&'l ()>,
}

impl<'l> StepWithGraph<'l> {
    /// Creates a StepWithGraph.
    pub fn new() -> Self {
        StepWithGraph {
            input_ports: vec![],
            input_tensors: vec![],

            output_ports: vec![],
            output_tensors: vec![],

            target_operations: vec![],

            phantom: marker::PhantomData,
        }
    }

    /// Adds an input to be fed to the graph.
    pub fn add_input<T: TensorType>(&mut self,
                                    operation: &Operation,
                                    index: c_int,
                                    tensor: &'l Tensor<T>) {
        self.input_ports.push(tf::TF_Output {
                                  oper: operation.inner(),
                                  index: index,
                              });
        self.input_tensors.push(tensor.inner);
    }

    /// Requests that an output is fetched from the graph after running this step.
    /// Returns an index that you can then use to fetch this output from the step after running it.
    pub fn request_output(&mut self, operation: &Operation, index: c_int) -> OutputToken {
        self.output_ports.push(tf::TF_Output {
                                   oper: operation.inner(),
                                   index: index,
                               });
        self.output_tensors.push(ptr::null_mut());
        OutputToken { index: self.output_tensors.len() - 1 }
    }

    /// Extracts a tensor output given an index. A given index can only be extracted
    /// once per `Session::run`.
    /// Returns an error if output_idx is out of range, output is unavailable or the
    /// requested type does not match the type of the actual tensor.
    pub fn take_output<T: TensorType>(&mut self, token: OutputToken) -> Result<Tensor<T>> {
        let output_idx = token.index;
        if output_idx >= self.output_tensors.len() {
            return Err(Status::new_set(Code::OutOfRange,
                                       &format!("Requested output index is out of range: {} vs \
                                                 {}",
                                                output_idx,
                                                self.output_tensors.len()))
                               .unwrap());
        }
        if self.output_tensors[output_idx].is_null() {
            return Err(Status::new_set(Code::Unavailable,
                                       "Output not available. Either it was already taken, or \
                                        this step has not been sucessfully run yet.")
                               .unwrap());
        }
        let actual_data_type = self.output_data_type(output_idx).unwrap();
        if actual_data_type != T::data_type() {
            return Err(invalid_arg!("Requested tensor type does not match actual tensor type: \
                                     {} vs {}",
                                    actual_data_type,
                                    T::data_type()));
        }
        let tensor = unsafe { Tensor::from_tf_tensor(self.output_tensors[output_idx]).unwrap() };
        self.output_tensors[output_idx] = ptr::null_mut();
        Ok(tensor)
    }

    /// Adds a target operation to be executed when running the graph.
    pub fn add_target(&mut self, operation: &Operation) {
        self.target_operations.push(operation.inner());
    }

    /// Retuns the type of the tensor given an index.
    /// Returns `None` if the index is out of range or the output is not yet available.
    pub fn output_data_type(&self, output_idx: usize) -> Option<DataType> {
        if output_idx >= self.output_tensors.len() {
            return None;
        }
        if self.output_tensors[output_idx].is_null() {
            return None;
        }
        unsafe { Some(DataType::from_c(tf::TF_TensorType(self.output_tensors[output_idx]))) }
    }

    fn drop_output_tensors(&mut self) {
        for mut tensor in &mut self.output_tensors {
            // TODO: Is TF_DeleteTensor NULL safe?
            if !tensor.is_null() {
                unsafe {
                    tf::TF_DeleteTensor(*tensor);
                }
            }
            *tensor = ptr::null_mut();
        }
    }
}

impl<'l> Drop for StepWithGraph<'l> {
    fn drop(&mut self) {
        self.drop_output_tensors();
    }
}

////////////////////////

#[cfg(test)]
mod tests {
    use super::*;
    use super::super::DataType;
    use super::super::Graph;
    use super::super::Operation;
    use super::super::Output;
    use super::super::SessionOptions;
    use super::super::Shape;
    use super::super::Tensor;

    fn create_session() -> (Session, Operation, Operation) {
        let mut g = Graph::new();
        let two = {
            let mut nd = g.new_operation("Const", "two").unwrap();
            nd.set_attr_type("dtype", DataType::Float).unwrap();
            let mut value = Tensor::new(&[1]);
            value[0] = 2.0f32;
            nd.set_attr_tensor("value", value).unwrap();
            nd.finish().unwrap()
        };
        let x = {
            let mut nd = g.new_operation("Placeholder", "x").unwrap();
            nd.set_attr_type("dtype", DataType::Float).unwrap();
            nd.set_attr_shape("shape", &Shape(Some(vec![]))).unwrap();
            nd.finish().unwrap()
        };
        let y = {
            let mut nd = g.new_operation("Mul", "y").unwrap();
            nd.add_input(Output {
                             operation: two,
                             index: 0,
                         });
            nd.add_input(Output {
                             operation: x.clone(),
                             index: 0,
                         });
            nd.finish().unwrap()
        };
        let options = SessionOptions::new();
        match Session::new(&options, &g) {
            Ok(session) => (session, x, y),
            Err(status) => panic!("Creating session failed with status: {}", status),
        }
    }

    #[test]
    fn smoke() {
        create_session();
    }

    #[test]
    fn test_close() {
        let (mut session, _, _) = create_session();
        let status = session.close();
        assert!(status.is_ok());
    }

    #[test]
    fn test_run() {
        let (mut session, x_operation, y_operation) = create_session();
        let mut x = <Tensor<f32>>::new(&[2]);
        x[0] = 2.0;
        x[1] = 3.0;
        let mut step = StepWithGraph::new();
        step.add_input(&x_operation, 0, &x);
        let output_token = step.request_output(&y_operation, 0);
        session.run(&mut step).unwrap();
        let output_tensor = step.take_output::<f32>(output_token).unwrap();
        assert_eq!(output_tensor.len(), 2);
        assert_eq!(output_tensor[0], 4.0);
        assert_eq!(output_tensor[1], 6.0);
    }
}
<|endoftext|>


<commit_before>//
// gash.rs
//
// Starting code for PS2
// Running on Rust 0.9
//
// University of Virginia - cs4414 Spring 2014
// Weilin Xu, David Evans
// Version 0.4
//

extern mod extra;

use std::libc::funcs::posix88;
use std::io::signal::{Listener, Interrupt};
use std::{io, run, os, str};
use std::io::buffered::BufferedReader;
use std::io::stdin;
use std::io::fs::File;
use extra::getopts;

struct Shell {
    cmd_prompt: ~str,
    history: ~[~str],
}

impl Shell {
    fn new(prompt_str: &str) -> Shell {
        Shell {
            cmd_prompt: prompt_str.to_owned(),
            history: ~[],
        }
    }

    fn run(&mut self) {
        let mut stdin = BufferedReader::new(stdin());

        loop {
            let cwd = std::os::getcwd();
            let cwd2 = format!("{}", cwd.display());
            let mut cwdS : ~[&str] = cwd2.split('/').collect();
            print!("{} : {}", cwdS.pop(), self.cmd_prompt);
            io::stdio::flush();

            let line = stdin.read_line().unwrap();
            let cmd_line = line.trim().to_owned();
            let writeRedirect = cmd_line.find_str(" > ");
            let readRedirect = cmd_line.find_str(" < ");
            if(writeRedirect == None && readRedirect == None) {
            	let program = cmd_line.splitn(' ', 1).nth(0).expect("no program");
            	self.history.push(cmd_line.clone());
            	match cmd_line.slice_from(cmd_line.len() - 2) {
            		" &"    => { self.run_background(program, cmd_line.slice(0, cmd_line.len() - 2));
            			continue; }
                	_	=> { }
            	}
            	match program {
			""           =>  { continue; }
			"exit"       =>  { return; }
			"cd"	     =>  { self.run_cd(cmd_line); }
			"history"    =>  { self.run_history(); }
			_            =>  { self.run_cmdline(cmd_line); }
            	}
            }
            else if (readRedirect == None) {
            	let cmd1 = cmd_line.slice(0, writeRedirect.unwrap());
            	let cmd2 = cmd_line.slice_from(writeRedirect.unwrap()+3);
            	let program = cmd1.splitn(' ', 1).nth(0).expect("no program");
            	let args : ~[&str] = if (cmd1.splitn(' ', 1).nth(1) != None) {cmd1.splitn(' ', 1).nth(1).unwrap().split(' ').collect()} else {~[]};
            	let mut argsOwned : ~[~str] = ~[];
            	for i in range (0, args.len()) {
            		argsOwned.push(args[i].to_owned());
            	}
		let newStdOut = File::create(&Path::new(cmd2));
            	match newStdOut {
            		Some(mut x) => {
            			let process = run::Process::new(program, argsOwned, run::ProcessOptions::new());
            			let output = process.unwrap().finish_with_output();
            			x.write(output.output);
            		}
            		None => { println("shouldnt be here");}
            	}
            }
            else if (writeRedirect == None) {
            	let cmd1 = cmd_line.slice(0, readRedirect.unwrap());
            	let cmd2 = cmd_line.slice_from(readRedirect.unwrap()+3);
            	let program = cmd1.splitn(' ', 1).nth(0).expect("no program");
            	let args : ~[&str] = if (cmd1.splitn(' ', 1).nth(1) != None) {cmd1.splitn(' ', 1).nth(1).unwrap().split(' ').collect()} else {~[]};
            	let mut argsOwned : ~[~str] = ~[];
            	for i in range (0, args.len()) {
            		argsOwned.push(args[i].to_owned());
            	}
           	let path = &Path::new(cmd2);
            	if (path.exists()) {
		let input = File::open(path);
            		match input {
            			Some(mut x) => {
            				let inputBytes = x.read_to_end();
            				let process = run::Process::new(program, argsOwned, run::ProcessOptions::new());
            				match(process) {
            					Some(mut toWrite) => {
            						toWrite.input().write(inputBytes);
            						let output = toWrite.finish_with_output();
            						print!("{:s}", str::from_utf8(output.output));
            					}
            					None => {}
        				}
            			}
            			None => { println!("gash: {:s}: No such file or directory", cmd2);}
            		}
            	}
                else {
           		println!("gash: {:s}: No such file or directory", cmd2);
           	}
            }
        }
    }

    fn run_cmdline(&mut self, cmd_line: &str) {
        let mut argv: ~[~str] =
            cmd_line.split(' ').filter_map(|x| if x != "" { Some(x.to_owned()) } else { None }).to_owned_vec();

        if argv.len() > 0 {
            let program: ~str = argv.remove(0);
            self.run_cmd(program, argv);
        }
    }

    fn run_cmd(&mut self, program: &str, argv: &[~str]) {
        if self.cmd_exists(program) {
            run::process_status(program, argv);
        } else {
            println!("{:s}: command not found", program);
        }
    }

    fn run_cd(&mut self, program: &str) {
        let mut argv: ~[~str] =
            program.split(' ').filter_map(|x| if x != "" { Some(x.to_owned()) } else { None }).to_owned_vec();

    	let mut programs : ~str = ~"";
            if argv.len() > 1 {
                programs = argv.remove(1);
            }
    	os::change_dir(&Path::new(programs.clone()));
    }

    fn run_history(&mut self) {
        for c in range(0, self.history.len()) {
            println!("{:s}", self.history[c]);
        }
    }

    fn run_background(&mut self, program: &str, cmd_line: &str) {
		let (portSelf, chanSelf): (Port<~Shell>, Chan<~Shell>) = Chan::new();
		let (portProg, chanProg): (Port<~str>, Chan<~str>) = Chan::new();
		let (portLine, chanLine): (Port<~str>, Chan<~str>) = Chan::new();

		let mut sendShell = ~(Shell::new("gash > "));
		sendShell.history = self.history.clone();
		chanSelf.send(sendShell);
		chanProg.send(program.to_owned());
		chanLine.send(cmd_line.to_owned());

		do spawn {
		let program : ~str = portProg.recv();
		let cmd_line : ~str = portLine.recv();
		let mut selfV : ~Shell = portSelf.recv();
                match program {
             	   ~""           =>  { }
              	  ~"exit"       =>  { }
		  ~"cd"	     =>  { selfV.run_cd(cmd_line); }
                  ~"history"    =>  { selfV.run_history(); }
                  _            =>  { selfV.run_cmdline(cmd_line); }
                }
            };
    }

    fn cmd_exists(&mut self, cmd_path: &str) -> bool {
        let ret = run::process_output("which", [cmd_path.to_owned()]);
        return ret.expect("exit code error.").status.success();
    }
}

fn get_cmdline_from_args() -> Option<~str> {
    /* Begin processing program arguments and initiate the parameters. */
    let args = os::args();

    let opts = ~[
        getopts::optopt("c")
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => { m }
        Err(f) => { fail!(f.to_err_msg()) }
    };

    if matches.opt_present("c") {
        let cmd_str = match matches.opt_str("c") {
                                                Some(cmd_str) => {cmd_str.to_owned()},
                                                None => {~""}
                                              };
        return Some(cmd_str);
    } else {
        return None;
    }
}

fn main() {
    let opt_cmd_line = get_cmdline_from_args();

    let mut listener = Listener::new();
    listener.register(Interrupt);

    let (portSelf, chanSelf): (Port<Listener>, Chan<Listener>) = Chan::new();
    chanSelf.send(listener);

    do spawn{
        loop {
            let listener = portSelf.recv();
            match listener.port.recv() {
                Interrupt => { unsafe { posix88::signal::kill(std::libc::getpid() , std::libc::SIGINT); }
                            }
                _ => (),
            }
        }
    };

    match opt_cmd_line {
        Some(cmd_line) => Shell::new("").run_cmdline(cmd_line),
        None           => Shell::new("gash > ").run()
    }
}
<commit_msg>running commands from history - problem 6<commit_after>//
// gash.rs
//
// Starting code for PS2
// Running on Rust 0.9
//
// University of Virginia - cs4414 Spring 2014
// Weilin Xu, David Evans
// Version 0.4
//

extern mod extra;

use std::libc::funcs::posix88;
use std::io::signal::{Listener, Interrupt};
use std::{io, run, os, str};
use std::io::buffered::BufferedReader;
use std::io::stdin;
use std::io::fs::File;
use extra::getopts;

struct Shell {
    cmd_prompt: ~str,
    history: ~[~str],
}

impl Shell {
    fn new(prompt_str: &str) -> Shell {
        Shell {
            cmd_prompt: prompt_str.to_owned(),
            history: ~[],
        }
    }

    fn run(&mut self) {
        let mut stdin = BufferedReader::new(stdin());

        loop {
            let cwd = std::os::getcwd();
            let cwd2 = format!("{}", cwd.display());
            let mut cwdS : ~[&str] = cwd2.split('/').collect();
            let last = cwdS.pop();
            if(last.len() > 0) {
                print!("{} : {}", last, self.cmd_prompt);
            } else {
                print!("Home : {}", self.cmd_prompt);
            }
            io::stdio::flush();

            let line = stdin.read_line().unwrap();
            let cmd_line = line.trim().to_owned();
            let writeRedirect = cmd_line.find_str(" > ");
            let readRedirect = cmd_line.find_str(" < ");
            if(writeRedirect == None && readRedirect == None) {
            	let program = cmd_line.splitn(' ', 1).nth(0).expect("no program");
            	self.history.push(cmd_line.clone());
            	match cmd_line.slice_from(cmd_line.len() - 2) {
            		" &"    => { self.run_background(program, cmd_line.slice(0, cmd_line.len() - 2));
            			continue; }
                	_	=> { }
            	}
            	match program {
			""           =>  { continue; }
			"exit"       =>  { return; }
			"cd"	     =>  { self.run_cd(cmd_line); }
			"history"    =>  { self.run_history(cmd_line); }
			_            =>  { self.run_cmdline(cmd_line); }
            	}
            }
            else if (readRedirect == None) {
            	let cmd1 = cmd_line.slice(0, writeRedirect.unwrap());
            	let cmd2 = cmd_line.slice_from(writeRedirect.unwrap()+3);
            	let program = cmd1.splitn(' ', 1).nth(0).expect("no program");
            	let args : ~[&str] = if (cmd1.splitn(' ', 1).nth(1) != None) {cmd1.splitn(' ', 1).nth(1).unwrap().split(' ').collect()} else {~[]};
            	let mut argsOwned : ~[~str] = ~[];
            	for i in range (0, args.len()) {
            		argsOwned.push(args[i].to_owned());
            	}
		      let newStdOut = File::create(&Path::new(cmd2));
            	match newStdOut {
            		Some(mut x) => {
            			let process = run::Process::new(program, argsOwned, run::ProcessOptions::new());
            			let output = process.unwrap().finish_with_output();
            			x.write(output.output);
            		}
            		None => { println("shouldnt be here");}
            	}
            }
            else if (writeRedirect == None) {
            	let cmd1 = cmd_line.slice(0, readRedirect.unwrap());
            	let cmd2 = cmd_line.slice_from(readRedirect.unwrap()+3);
            	let program = cmd1.splitn(' ', 1).nth(0).expect("no program");
            	let args : ~[&str] = if (cmd1.splitn(' ', 1).nth(1) != None) {cmd1.splitn(' ', 1).nth(1).unwrap().split(' ').collect()} else {~[]};
            	let mut argsOwned : ~[~str] = ~[];
            	for i in range (0, args.len()) {
            		argsOwned.push(args[i].to_owned());
            	}
           	let path = &Path::new(cmd2);
            	if (path.exists()) {
		let input = File::open(path);
            		match input {
            			Some(mut x) => {
            				let inputBytes = x.read_to_end();
            				let process = run::Process::new(program, argsOwned, run::ProcessOptions::new());
            				match(process) {
            					Some(mut toWrite) => {
            						toWrite.input().write(inputBytes);
            						let output = toWrite.finish_with_output();
            						print!("{:s}", str::from_utf8(output.output));
            					}
            					None => {}
        				}
            			}
            			None => { println!("gash: {:s}: No such file or directory", cmd2);}
            		}
            	}
                else {
           		println!("gash: {:s}: No such file or directory", cmd2);
           	}
            }
        }
    }

    fn run_cmdline(&mut self, cmd_line: &str) {
        let mut argv: ~[~str] =
            cmd_line.split(' ').filter_map(|x| if x != "" { Some(x.to_owned()) } else { None }).to_owned_vec();

        if argv.len() > 0 {
            let program: ~str = argv.remove(0);
            self.run_cmd(program, argv);
        }
    }

    fn run_cmd(&mut self, program: &str, argv: &[~str]) {
        if self.cmd_exists(program) {
            run::process_status(program, argv);
        } else {
            println!("{:s}: command not found", program);
        }
    }

    fn run_cd(&mut self, program: &str) {
        let mut argv: ~[~str] =
            program.split(' ').filter_map(|x| if x != "" { Some(x.to_owned()) } else { None }).to_owned_vec();

    	let mut programs : ~str = ~"";
            if argv.len() > 1 {
                programs = argv.remove(1);
            }
    	os::change_dir(&Path::new(programs.clone()));
    }

    fn run_history(&mut self, program: &str) {
        let histArgs : ~[&str] = program.split(' ').collect();
        if(histArgs.len() == 2) {
            if(from_str::<uint>(histArgs[1]).unwrap() < self.history.len()) {
                let num = self.history.len() - from_str::<uint>(histArgs[1]).unwrap() - 1;
                println!("command to run: {:s}", self.history[num]);
                let cmd = self.history[num].to_owned();
                self.history.push(cmd);
                // run self.history[num] exactly as printed
            } else {
                println("You haven't entered that many commands yet - try a smaller number");
            }
        }
        for c in range(0, self.history.len()) {
            println!("{:s}", self.history[c]);
        }
    }

    fn run_background(&mut self, program: &str, cmd_line: &str) {
		let (portSelf, chanSelf): (Port<~Shell>, Chan<~Shell>) = Chan::new();
		let (portProg, chanProg): (Port<~str>, Chan<~str>) = Chan::new();
		let (portLine, chanLine): (Port<~str>, Chan<~str>) = Chan::new();

		let mut sendShell = ~(Shell::new("gash > "));
		sendShell.history = self.history.clone();
		chanSelf.send(sendShell);
		chanProg.send(program.to_owned());
		chanLine.send(cmd_line.to_owned());

		do spawn {
		let program : ~str = portProg.recv();
		let cmd_line : ~str = portLine.recv();
		let mut selfV : ~Shell = portSelf.recv();
                match program {
             	   ~""           =>  { }
              	  ~"exit"       =>  { }
		  ~"cd"	     =>  { selfV.run_cd(cmd_line); }
                  ~"history"    =>  { selfV.run_history(cmd_line); }
                  _            =>  { selfV.run_cmdline(cmd_line); }
                }
            };
    }

    fn cmd_exists(&mut self, cmd_path: &str) -> bool {
        let ret = run::process_output("which", [cmd_path.to_owned()]);
        return ret.expect("exit code error.").status.success();
    }
}

fn get_cmdline_from_args() -> Option<~str> {
    /* Begin processing program arguments and initiate the parameters. */
    let args = os::args();

    let opts = ~[
        getopts::optopt("c")
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => { m }
        Err(f) => { fail!(f.to_err_msg()) }
    };

    if matches.opt_present("c") {
        let cmd_str = match matches.opt_str("c") {
                                                Some(cmd_str) => {cmd_str.to_owned()},
                                                None => {~""}
                                              };
        return Some(cmd_str);
    } else {
        return None;
    }
}

fn main() {
    let opt_cmd_line = get_cmdline_from_args();

    let mut listener = Listener::new();
    listener.register(Interrupt);

    /*let (portSelf, chanSelf): (Port<Listener>, Chan<Listener>) = Chan::new();
    chanSelf.send(listener);

    do spawn{
        loop {
            let listener = portSelf.recv();
            match listener.port.recv() {
                Interrupt => { unsafe { posix88::signal::kill(std::libc::getpid() , std::libc::SIGINT); }
                            }
                _ => (),
            }
        }
    };*/

    match opt_cmd_line {
        Some(cmd_line) => Shell::new("").run_cmdline(cmd_line),
        None           => Shell::new("gash > ").run()
    }
}
<|endoftext|>


<commit_before>//
// gash.rs
//
// Starting code for PS2
// Running on Rust 0.9
//
// University of Virginia - cs4414 Spring 2014
// Weilin Xu, David Evans
// Version 0.4
//

extern mod extra;

use std::{io, run, os, str};
use std::io::buffered::BufferedReader;
use std::io::stdin;
use std::io::fs::File;
use extra::getopts;

struct Shell {
    cmd_prompt: ~str,
    history: ~[~str],
}

impl Shell {
    fn new(prompt_str: &str) -> Shell {
        Shell {
            cmd_prompt: prompt_str.to_owned(),
            history: ~[],
        }
    }

    fn run(&mut self) {
        let mut stdin = BufferedReader::new(stdin());

        loop {
            print(self.cmd_prompt);
            io::stdio::flush();

            let line = stdin.read_line().unwrap();
            let cmd_line = line.trim().to_owned();
            let writeRedirect = cmd_line.find_str(" > ");
            let readRedirect = cmd_line.find_str(" < ");
            if(writeRedirect == None && readRedirect == None) {
            	let program = cmd_line.splitn(' ', 1).nth(0).expect("no program");
            	self.history.push(cmd_line.clone());
            	match cmd_line.slice_from(cmd_line.len() - 2) {
            		" &"    => { self.run_background(program, cmd_line);
            			continue; }
                	_	=> { }
            	}
            	match program {
			""           =>  { continue; }
			"exit"       =>  { return; }
			"cd"	     =>  { self.run_cd(cmd_line); }
			"history"    =>  { self.run_history(); }
			_            =>  { self.run_cmdline(cmd_line); }
            	}
            }
            else if (readRedirect == None) {
            	let cmd1 = cmd_line.slice(0, writeRedirect.unwrap());
            	let cmd2 = cmd_line.slice_from(writeRedirect.unwrap()+3);
            	let program = cmd1.splitn(' ', 1).nth(0).expect("no program");
            	let args : ~[&str] = if (cmd1.splitn(' ', 1).nth(1) != None) {cmd1.splitn(' ', 1).nth(1).unwrap().split(' ').collect()} else {~[]};
            	let mut argsOwned : ~[~str] = ~[];
            	for i in range (0, args.len()) {
            		argsOwned.push(args[i].to_owned());
            	}
		let newStdOut = File::create(&Path::new(cmd2));
            	match newStdOut {
            		Some(mut x) => {
            			let process = run::Process::new(program, argsOwned, run::ProcessOptions::new());
            			let output = process.unwrap().finish_with_output();
            			x.write(output.output);
            		}
            		None => { println("shouldnt be here");}
            	}
            }
            else if (writeRedirect == None) {
            	let cmd1 = cmd_line.slice(0, readRedirect.unwrap());
            	let cmd2 = cmd_line.slice_from(readRedirect.unwrap()+3);
            	let program = cmd1.splitn(' ', 1).nth(0).expect("no program");
            	let args : ~[&str] = if (cmd1.splitn(' ', 1).nth(1) != None) {cmd1.splitn(' ', 1).nth(1).unwrap().split(' ').collect()} else {~[]};
            	let mut argsOwned : ~[~str] = ~[];
            	for i in range (0, args.len()) {
            		argsOwned.push(args[i].to_owned());
            	}
		let input = File::open(&Path::new(cmd2));
            	match input {
            		Some(mut x) => {
            			let inputBytes = x.read_to_end();
            			let process = run::Process::new(program, argsOwned, run::ProcessOptions::new());
            			match(process) {
            				Some(mut toWrite) => {
            					toWrite.input().write(inputBytes);
            					let output = toWrite.finish_with_output();
            					print!("{:s}", str::from_utf8(output.output));
            				}
            				None => {}
        			}
            		}
            		None => { println!("gash: {:s}: No such file or directory", cmd2);}
            	}
            }
        }
    }

    fn run_cmdline(&mut self, cmd_line: &str) {
        let mut argv: ~[~str] =
            cmd_line.split(' ').filter_map(|x| if x != "" { Some(x.to_owned()) } else { None }).to_owned_vec();

        if argv.len() > 0 {
            let program: ~str = argv.remove(0);
            self.run_cmd(program, argv);
        }
    }

    fn run_cmd(&mut self, program: &str, argv: &[~str]) {
        if self.cmd_exists(program) {
            run::process_status(program, argv);
        } else {
            println!("{:s}: command not found", program);
        }
    }

    fn run_cd(&mut self, program: &str) {
        let mut argv: ~[~str] =
            program.split(' ').filter_map(|x| if x != "" { Some(x.to_owned()) } else { None }).to_owned_vec();

    	let mut programs : ~str = ~"";
            if argv.len() > 0 {
                programs = argv.remove(1);
            }
    	os::change_dir(&Path::new(programs.clone()));
    }

    fn run_history(&mut self) {
        for c in range(0, self.history.len()) {
            println!("{:s}", self.history[c]);
        }
    }

    fn run_background(&mut self, program: &str, cmd_line: &str) {
		let (portSelf, chanSelf): (Port<~Shell>, Chan<~Shell>) = Chan::new();
		let (portProg, chanProg): (Port<~str>, Chan<~str>) = Chan::new();
		let (portLine, chanLine): (Port<~str>, Chan<~str>) = Chan::new();

		chanSelf.send(~(Shell::new("gash > ")));
		chanProg.send(program.to_owned());
		chanLine.send(cmd_line.to_owned());
            do spawn {
		let program : ~str = portProg.recv();
		let cmd_line : ~str = portLine.recv();
		let mut selfV : ~Shell = portSelf.recv();
                match program {
             	   ~""           =>  { }
              	  ~"exit"       =>  { }
		  ~"cd"	     =>  { selfV.run_cd(cmd_line); }
                  ~"history"    =>  { selfV.run_history(); }
                  _            =>  { selfV.run_cmdline(cmd_line); }
                }
            };
    }

    fn cmd_exists(&mut self, cmd_path: &str) -> bool {
        let ret = run::process_output("which", [cmd_path.to_owned()]);
        return ret.expect("exit code error.").status.success();
    }
}

fn get_cmdline_from_args() -> Option<~str> {
    /* Begin processing program arguments and initiate the parameters. */
    let args = os::args();

    let opts = ~[
        getopts::optopt("c")
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => { m }
        Err(f) => { fail!(f.to_err_msg()) }
    };

    if matches.opt_present("c") {
        let cmd_str = match matches.opt_str("c") {
                                                Some(cmd_str) => {cmd_str.to_owned()},
                                                None => {~""}
                                              };
        return Some(cmd_str);
    } else {
        return None;
    }
}

fn main() {
    let opt_cmd_line = get_cmdline_from_args();

    match opt_cmd_line {
        Some(cmd_line) => Shell::new("").run_cmdline(cmd_line),
        None           => Shell::new("gash > ").run()
    }
}
<commit_msg>< works even if file is not there<commit_after>//
// gash.rs
//
// Starting code for PS2
// Running on Rust 0.9
//
// University of Virginia - cs4414 Spring 2014
// Weilin Xu, David Evans
// Version 0.4
//

extern mod extra;

use std::{io, run, os, str};
use std::io::buffered::BufferedReader;
use std::io::stdin;
use std::io::fs::File;
use extra::getopts;

struct Shell {
    cmd_prompt: ~str,
    history: ~[~str],
}

impl Shell {
    fn new(prompt_str: &str) -> Shell {
        Shell {
            cmd_prompt: prompt_str.to_owned(),
            history: ~[],
        }
    }

    fn run(&mut self) {
        let mut stdin = BufferedReader::new(stdin());

        loop {
            print(self.cmd_prompt);
            io::stdio::flush();

            let line = stdin.read_line().unwrap();
            let cmd_line = line.trim().to_owned();
            let writeRedirect = cmd_line.find_str(" > ");
            let readRedirect = cmd_line.find_str(" < ");
            if(writeRedirect == None && readRedirect == None) {
            	let program = cmd_line.splitn(' ', 1).nth(0).expect("no program");
            	self.history.push(cmd_line.clone());
            	match cmd_line.slice_from(cmd_line.len() - 2) {
            		" &"    => { self.run_background(program, cmd_line);
            			continue; }
                	_	=> { }
            	}
            	match program {
			""           =>  { continue; }
			"exit"       =>  { return; }
			"cd"	     =>  { self.run_cd(cmd_line); }
			"history"    =>  { self.run_history(); }
			_            =>  { self.run_cmdline(cmd_line); }
            	}
            }
            else if (readRedirect == None) {
            	let cmd1 = cmd_line.slice(0, writeRedirect.unwrap());
            	let cmd2 = cmd_line.slice_from(writeRedirect.unwrap()+3);
            	let program = cmd1.splitn(' ', 1).nth(0).expect("no program");
            	let args : ~[&str] = if (cmd1.splitn(' ', 1).nth(1) != None) {cmd1.splitn(' ', 1).nth(1).unwrap().split(' ').collect()} else {~[]};
            	let mut argsOwned : ~[~str] = ~[];
            	for i in range (0, args.len()) {
            		argsOwned.push(args[i].to_owned());
            	}
		let newStdOut = File::create(&Path::new(cmd2));
            	match newStdOut {
            		Some(mut x) => {
            			let process = run::Process::new(program, argsOwned, run::ProcessOptions::new());
            			let output = process.unwrap().finish_with_output();
            			x.write(output.output);
            		}
            		None => { println("shouldnt be here");}
            	}
            }
            else if (writeRedirect == None) {
            	let cmd1 = cmd_line.slice(0, readRedirect.unwrap());
            	let cmd2 = cmd_line.slice_from(readRedirect.unwrap()+3);
            	let program = cmd1.splitn(' ', 1).nth(0).expect("no program");
            	let args : ~[&str] = if (cmd1.splitn(' ', 1).nth(1) != None) {cmd1.splitn(' ', 1).nth(1).unwrap().split(' ').collect()} else {~[]};
            	let mut argsOwned : ~[~str] = ~[];
            	for i in range (0, args.len()) {
            		argsOwned.push(args[i].to_owned());
            	}
           	let path = &Path::new(cmd2);
            	if (path.exists()) {
		let input = File::open(path);
            		match input {
            			Some(mut x) => {
            				let inputBytes = x.read_to_end();
            				let process = run::Process::new(program, argsOwned, run::ProcessOptions::new());
            				match(process) {
            					Some(mut toWrite) => {
            						toWrite.input().write(inputBytes);
            						let output = toWrite.finish_with_output();
            						print!("{:s}", str::from_utf8(output.output));
            					}
            					None => {}
        				}
            			}
            			None => { println!("gash: {:s}: No such file or directory", cmd2);}
            		}
            	}
                else {
           		println!("gash: {:s}: No such file or directory", cmd2);
           	}
            }
        }
    }

    fn run_cmdline(&mut self, cmd_line: &str) {
        let mut argv: ~[~str] =
            cmd_line.split(' ').filter_map(|x| if x != "" { Some(x.to_owned()) } else { None }).to_owned_vec();

        if argv.len() > 0 {
            let program: ~str = argv.remove(0);
            self.run_cmd(program, argv);
        }
    }

    fn run_cmd(&mut self, program: &str, argv: &[~str]) {
        if self.cmd_exists(program) {
            run::process_status(program, argv);
        } else {
            println!("{:s}: command not found", program);
        }
    }

    fn run_cd(&mut self, program: &str) {
        let mut argv: ~[~str] =
            program.split(' ').filter_map(|x| if x != "" { Some(x.to_owned()) } else { None }).to_owned_vec();

    	let mut programs : ~str = ~"";
            if argv.len() > 0 {
                programs = argv.remove(1);
            }
    	os::change_dir(&Path::new(programs.clone()));
    }

    fn run_history(&mut self) {
        for c in range(0, self.history.len()) {
            println!("{:s}", self.history[c]);
        }
    }

    fn run_background(&mut self, program: &str, cmd_line: &str) {
		let (portSelf, chanSelf): (Port<~Shell>, Chan<~Shell>) = Chan::new();
		let (portProg, chanProg): (Port<~str>, Chan<~str>) = Chan::new();
		let (portLine, chanLine): (Port<~str>, Chan<~str>) = Chan::new();

		chanSelf.send(~(Shell::new("gash > ")));
		chanProg.send(program.to_owned());
		chanLine.send(cmd_line.to_owned());
            do spawn {
		let program : ~str = portProg.recv();
		let cmd_line : ~str = portLine.recv();
		let mut selfV : ~Shell = portSelf.recv();
                match program {
             	   ~""           =>  { }
              	  ~"exit"       =>  { }
		  ~"cd"	     =>  { selfV.run_cd(cmd_line); }
                  ~"history"    =>  { selfV.run_history(); }
                  _            =>  { selfV.run_cmdline(cmd_line); }
                }
            };
    }

    fn cmd_exists(&mut self, cmd_path: &str) -> bool {
        let ret = run::process_output("which", [cmd_path.to_owned()]);
        return ret.expect("exit code error.").status.success();
    }
}

fn get_cmdline_from_args() -> Option<~str> {
    /* Begin processing program arguments and initiate the parameters. */
    let args = os::args();

    let opts = ~[
        getopts::optopt("c")
    ];

    let matches = match getopts::getopts(args.tail(), opts) {
        Ok(m) => { m }
        Err(f) => { fail!(f.to_err_msg()) }
    };

    if matches.opt_present("c") {
        let cmd_str = match matches.opt_str("c") {
                                                Some(cmd_str) => {cmd_str.to_owned()},
                                                None => {~""}
                                              };
        return Some(cmd_str);
    } else {
        return None;
    }
}

fn main() {
    let opt_cmd_line = get_cmdline_from_args();

    match opt_cmd_line {
        Some(cmd_line) => Shell::new("").run_cmdline(cmd_line),
        None           => Shell::new("gash > ").run()
    }
}
<|endoftext|>


<commit_before>use std::cell::RefCell;
use std::rc::Rc;

use url::Url;
use hyper::{self, Client};
use hyper::method::Method;
use rustc_serialize::json::{self, Json, Object};

use errors::WeChatError;

//mod misc;

#[derive(Debug, Eq, PartialEq, Clone)]
pub struct WeChatClient {
    pub appid: String,
    pub secret: String,
    _access_token: Rc<RefCell<String>>,
}

impl WeChatClient {

    #[inline]
    pub fn new(appid: &str, secret: &str) -> WeChatClient {
        WeChatClient {
            appid: appid.to_owned(),
            secret: secret.to_owned(),
            _access_token: Rc::new(RefCell::new("".to_owned())),
        }
    }

    #[inline]
    pub fn with_access_token(appid: &str, secret: &str, access_token: &str) -> WeChatClient {
        WeChatClient {
            appid: appid.to_owned(),
            secret: secret.to_owned(),
            _access_token: Rc::new(RefCell::new(access_token.to_owned())),
        }
    }

    #[inline]
    pub fn access_token(&self) -> String {
        self._access_token.borrow().clone()
    }

    pub fn request(&self, method: Method, url: &str, params: Vec<(&str, &str)>, data: &Object) -> Result<Json, WeChatError> {
        let mut http_url = if !(url.starts_with("http://") || url.starts_with("https://")) {
            let mut url_string = "https://api.weixin.qq.com/cgi-bin/".to_owned();
            url_string = url_string + url;
            Url::parse(&url_string).unwrap()
        } else {
            Url::parse(url).unwrap()
        };
        let access_token = self.access_token();
        let mut querys = params.clone();
        if !access_token.is_empty() {
            debug!("Using access_token: {}", access_token);
            querys.push(("access_token", &access_token));
        }
        http_url.set_query_from_pairs(querys.into_iter());
        let body = if !data.is_empty() {
            match json::encode(data) {
                Ok(text) => text,
                Err(_) => "".to_owned(),
            }
        } else {
            "".to_owned()
        };
        let mut client = Client::new();
        let req = if method == Method::Post {
            client.post(http_url).body(&body)
        } else {
            client.get(http_url)
        };
        let mut res = match req.send() {
            Ok(_res) => _res,
            Err(_) => {
                error!("Send request error");
                return Err(WeChatError::ClientError { errcode: -1, errmsg: "Send request error".to_owned() });
            }
        };
        if res.status != hyper::Ok {
            error!("Response status error");
            return Err(WeChatError::ClientError { errcode: -2, errmsg: "Response status error".to_owned() })
        }
        let obj = match Json::from_reader(&mut res) {
            Ok(decoded) => { decoded },
            Err(_) => {
                error!("Json decode error");
                return Err(WeChatError::ClientError { errcode: -3, errmsg: "Json decode error".to_owned() });
            }
        };
        Ok(obj)
    }

    #[inline]
    pub fn post(&self, url: &str, params: Vec<(&str, &str)>, data: &Object) -> Result<Json, WeChatError> {
        if self.access_token().is_empty() {
            self.fetch_access_token();
        }
        self.request(Method::Post, url, params, data)
    }

    #[inline]
    pub fn get(&self, url: &str, params: Vec<(&str, &str)>) -> Result<Json, WeChatError> {
        if self.access_token().is_empty() {
            self.fetch_access_token();
        }
        self.request(Method::Get, url, params, &Object::new())
    }

    pub fn fetch_access_token(&self) -> Option<String> {
        let res = self.request(
            Method::Get,
            "token",
            vec![
                ("grant_type", "client_credential"),
                ("appid", &self.appid),
                ("secret", &self.secret),
            ],
            &Object::new()
        );
        let data = match res {
            Ok(data) => data,
            Err(_) => { return None; },
        };
        let token = match data.find("access_token") {
            Some(token) => token,
            None => { return None; }
        };
        let token_str = match *token {
            Json::String(ref v) => {
                let mut access_token = self._access_token.borrow_mut();
                *access_token = v.to_owned();
                Some(format!("{}", v))
            },
            _ => None,
        };
        token_str
    }
}
<commit_msg>Use Mutex and Arc instead of simple Rc to ensure threadsafety.<commit_after>use std::cell::RefCell;
use std::sync::{Mutex, Arc};

use url::Url;
use hyper::{self, Client};
use hyper::method::Method;
use rustc_serialize::json::{self, Json, Object};

use errors::WeChatError;

//mod misc;

#[derive(Debug, Clone)]
pub struct WeChatClient {
    pub appid: String,
    pub secret: String,
    _access_token: Arc<Mutex<RefCell<String>>>,
}

impl WeChatClient {

    #[inline]
    pub fn new(appid: &str, secret: &str) -> WeChatClient {
        WeChatClient {
            appid: appid.to_owned(),
            secret: secret.to_owned(),
            _access_token: Arc::new(Mutex::new(RefCell::new("".to_owned()))),
        }
    }

    #[inline]
    pub fn with_access_token(appid: &str, secret: &str, access_token: &str) -> WeChatClient {
        WeChatClient {
            appid: appid.to_owned(),
            secret: secret.to_owned(),
            _access_token: Arc::new(Mutex::new(RefCell::new(access_token.to_owned()))),
        }
    }

    #[inline]
    pub fn access_token(&self) -> String {
        let data = self._access_token.clone();
        let data = data.lock().unwrap();
        let token = data.borrow().clone();
        token
    }

    pub fn request(&self, method: Method, url: &str, params: Vec<(&str, &str)>, data: &Object) -> Result<Json, WeChatError> {
        let mut http_url = if !(url.starts_with("http://") || url.starts_with("https://")) {
            let mut url_string = "https://api.weixin.qq.com/cgi-bin/".to_owned();
            url_string = url_string + url;
            Url::parse(&url_string).unwrap()
        } else {
            Url::parse(url).unwrap()
        };
        let access_token = self.access_token();
        let mut querys = params.clone();
        if !access_token.is_empty() {
            debug!("Using access_token: {}", access_token);
            querys.push(("access_token", &access_token));
        }
        http_url.set_query_from_pairs(querys.into_iter());
        let body = if !data.is_empty() {
            match json::encode(data) {
                Ok(text) => text,
                Err(_) => "".to_owned(),
            }
        } else {
            "".to_owned()
        };
        let mut client = Client::new();
        let req = if method == Method::Post {
            client.post(http_url).body(&body)
        } else {
            client.get(http_url)
        };
        let mut res = match req.send() {
            Ok(_res) => _res,
            Err(_) => {
                error!("Send request error");
                return Err(WeChatError::ClientError { errcode: -1, errmsg: "Send request error".to_owned() });
            }
        };
        if res.status != hyper::Ok {
            error!("Response status error");
            return Err(WeChatError::ClientError { errcode: -2, errmsg: "Response status error".to_owned() })
        }
        let obj = match Json::from_reader(&mut res) {
            Ok(decoded) => { decoded },
            Err(_) => {
                error!("Json decode error");
                return Err(WeChatError::ClientError { errcode: -3, errmsg: "Json decode error".to_owned() });
            }
        };
        Ok(obj)
    }

    #[inline]
    pub fn post(&self, url: &str, params: Vec<(&str, &str)>, data: &Object) -> Result<Json, WeChatError> {
        if self.access_token().is_empty() {
            self.fetch_access_token();
        }
        self.request(Method::Post, url, params, data)
    }

    #[inline]
    pub fn get(&self, url: &str, params: Vec<(&str, &str)>) -> Result<Json, WeChatError> {
        if self.access_token().is_empty() {
            self.fetch_access_token();
        }
        self.request(Method::Get, url, params, &Object::new())
    }

    pub fn fetch_access_token(&self) -> Option<String> {
        let res = self.request(
            Method::Get,
            "token",
            vec![
                ("grant_type", "client_credential"),
                ("appid", &self.appid),
                ("secret", &self.secret),
            ],
            &Object::new()
        );
        let data = match res {
            Ok(data) => data,
            Err(_) => { return None; },
        };
        let token = match data.find("access_token") {
            Some(token) => token,
            None => { return None; }
        };
        let token_str = match *token {
            Json::String(ref v) => {
                let data = self._access_token.clone();
                let data = data.lock().unwrap();
                let mut access_token = data.borrow_mut();
                *access_token = v.to_owned();
                Some(format!("{}", v))
            },
            _ => None,
        };
        token_str
    }
}
<|endoftext|>


<commit_before>use futures::{Async, Future, Poll, Sink};
use futures::stream::Stream;
use futures::sink::Send;
use futures_state_stream::{StateStream, StreamEvent};

use native_tls::TlsConnector;

use std::io;
use std::net::ToSocketAddrs;

use tokio::net::{ConnectFuture, TcpStream};
use tokio_io::AsyncRead;
use tokio_tls::{ConnectAsync, TlsConnectorExt};

use imap_proto::{Request, RequestId, State};
use imap_proto::builders::command::Command;
use proto::{ImapCodec, ImapTls, ResponseData};

pub mod builder {
    pub use imap_proto::builders::command::{CommandBuilder, FetchBuilderAttributes,
                                            FetchBuilderMessages, FetchBuilderModifiers,
                                            FetchCommand, FetchCommandAttributes,
                                            FetchCommandMessages};
}

pub struct Client {
    transport: ImapTls,
    state: ClientState,
}

impl Client {
    pub fn connect(server: &str) -> io::Result<ImapConnectFuture> {
        let addr = (server, 993).to_socket_addrs()?.next().ok_or_else(|| {
            io::Error::new(io::ErrorKind::Other, format!("no IP addresses found for {}", server))
        })?;
        Ok(ImapConnectFuture::TcpConnecting(TcpStream::connect(&addr), server.to_string()))
    }

    pub fn call(self, cmd: Command) -> ResponseStream {
        let Self { transport, mut state } = self;
        let request_id = state.request_ids.next().unwrap(); // safe: never returns Err
        let (cmd_bytes, next_state) = cmd.into_parts();
        let future = transport.send(Request(request_id.clone(), cmd_bytes));
        ResponseStream::new(future, state, request_id, next_state)
    }
}

pub struct ResponseStream {
    future: Option<Send<ImapTls>>,
    transport: Option<ImapTls>,
    state: Option<ClientState>,
    request_id: RequestId,
    next_state: Option<State>,
    done: bool,
}

impl ResponseStream {
    pub fn new(future: Send<ImapTls>, state: ClientState,
               request_id: RequestId, next_state: Option<State>) -> Self {
        Self {
            future: Some(future),
            transport: None,
            state: Some(state),
            request_id: request_id,
            next_state: next_state,
            done: false,
        }
    }
}

impl StateStream for ResponseStream {
    type Item = ResponseData;
    type State = Client;
    type Error = io::Error;
    fn poll(&mut self) -> Poll<StreamEvent<Self::Item, Self::State>, Self::Error> {
        match self.future.take() {
            Some(mut future) => match future.poll() {
                Ok(Async::Ready(transport)) => {
                    self.transport = Some(transport);
                },
                Ok(Async::NotReady) => {
                    self.future = Some(future);
                    return Ok(Async::NotReady);
                },
                Err(e) => {
                    return Err(e);
                },
            },
            None => {},
        }
        let mut transport = match self.transport.take() {
            None => return Ok(Async::NotReady),
            Some(mut transport) => transport,
        };
        if self.done {
            let mut state = self.state.take().unwrap(); // safe: initialized from start
            if let Some(next_state) = self.next_state.take() {
                state.state = next_state;
            }
            let client = Client { transport, state };
            return Ok(Async::Ready(StreamEvent::Done(client)));
        }
        match transport.poll() {
            Ok(Async::Ready(Some(rsp))) => {
                if let Some(req_id) = rsp.request_id() {
                    self.done = *req_id == self.request_id;
                };
                self.transport = Some(transport);
                return Ok(Async::Ready(StreamEvent::Next(rsp)));
            },
            Err(e) => {
                return Err(e);
            },
            _ => (),
        }
        self.transport = Some(transport);
        Ok(Async::NotReady)
    }
}

pub enum ImapConnectFuture {
    #[doc(hidden)] TcpConnecting(ConnectFuture, String),
    #[doc(hidden)] TlsHandshake(ConnectAsync<TcpStream>),
    #[doc(hidden)] ServerGreeting(Option<ImapTls>),
}

impl Future for ImapConnectFuture {
    type Item = (Client, ResponseData);
    type Error = io::Error;
    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let mut new = None;
        if let ImapConnectFuture::TcpConnecting(ref mut future, ref domain) = *self {
            let stream = try_ready!(future.poll());
            let ctx = TlsConnector::builder().unwrap().build().unwrap();
            let future = ctx.connect_async(domain, stream);
            new = Some(ImapConnectFuture::TlsHandshake(future));
        }
        if new.is_some() {
            *self = new.take().unwrap();
        }
        if let ImapConnectFuture::TlsHandshake(ref mut future) = *self {
            let transport = try_ready!(future.map_err(|e| {
                io::Error::new(io::ErrorKind::Other, e)
            }).poll()).framed(ImapCodec::default());
            new = Some(ImapConnectFuture::ServerGreeting(Some(transport)));
        }
        if new.is_some() {
            *self = new.take().unwrap();
        }
        if let ImapConnectFuture::ServerGreeting(ref mut wrapped) = *self {
            let msg = try_ready!(wrapped.as_mut().unwrap().poll()).unwrap();
            return Ok(Async::Ready((Client {
                transport: wrapped.take().unwrap(),
                state: ClientState::new(),
            }, msg)));
        }
        Ok(Async::NotReady)
    }
}

pub struct ClientState {
    state: State,
    request_ids: IdGenerator,
}

impl ClientState {
    pub fn new() -> Self {
        Self {
            state: State::NotAuthenticated,
            request_ids: IdGenerator::new(),
        }
    }
}

impl Default for ClientState {
    fn default() -> Self {
        Self::new()
    }
}

pub struct IdGenerator {
    next: u64,
}

impl IdGenerator {
    pub fn new() -> Self {
        Self { next: 0 }
    }
}

impl Default for IdGenerator {
    fn default() -> Self {
        Self::new()
    }
}

impl Iterator for IdGenerator {
    type Item = RequestId;
    fn next(&mut self) -> Option<Self::Item> {
        self.next += 1;
        Some(RequestId(format!("A{:04}", self.next % 10_000)))
    }
}
<commit_msg>Disconnect core client types from concrete transport type<commit_after>use futures::{Async, Future, Poll, Sink};
use futures::stream::Stream;
use futures::sink::Send;
use futures_state_stream::{StateStream, StreamEvent};

use native_tls::TlsConnector;

use std::io;
use std::net::ToSocketAddrs;

use tokio::net::{ConnectFuture, TcpStream};
use tokio_io::AsyncRead;
use tokio_tls::{ConnectAsync, TlsConnectorExt};

use imap_proto::{Request, RequestId, State};
use imap_proto::builders::command::Command;
use proto::{ImapCodec, ImapTls, ImapTransport, ResponseData};

pub mod builder {
    pub use imap_proto::builders::command::{CommandBuilder, FetchBuilderAttributes,
                                            FetchBuilderMessages, FetchBuilderModifiers,
                                            FetchCommand, FetchCommandAttributes,
                                            FetchCommandMessages};
}

pub fn connect(server: &str) -> io::Result<ImapConnectFuture> {
    let addr = (server, 993).to_socket_addrs()?.next().ok_or_else(|| {
        io::Error::new(io::ErrorKind::Other, format!("no IP addresses found for {}", server))
    })?;
    Ok(ImapConnectFuture::TcpConnecting(TcpStream::connect(&addr), server.to_string()))
}

pub trait ImapClient {
    type Transport: ImapTransport;
    fn into_parts(self) -> (Self::Transport, ClientState);
    fn rebuild(transport: Self::Transport, state: ClientState) -> Self;

    fn call(self, cmd: Command) -> ResponseStream<Self> where Self: ImapClient + Sized {
        let (transport, mut state) = self.into_parts();
        let request_id = state.request_ids.next().unwrap(); // safe: never returns Err
        let (cmd_bytes, next_state) = cmd.into_parts();
        let future = transport.send(Request(request_id.clone(), cmd_bytes));
        ResponseStream::new(future, state, request_id, next_state)
    }
}

pub struct Client {
    transport: ImapTls,
    state: ClientState,
}

impl ImapClient for Client {
    type Transport = ImapTls;

    fn into_parts(self) -> (ImapTls, ClientState) {
        let Self { transport, state } = self;
        (transport, state)
    }

    fn rebuild(transport: ImapTls, state: ClientState) -> Client {
        Client { transport, state }
    }
}

pub struct ResponseStream<E> where E: ImapClient {
    future: Option<Send<E::Transport>>,
    transport: Option<E::Transport>,
    state: Option<ClientState>,
    request_id: RequestId,
    next_state: Option<State>,
    done: bool,
}

impl<E> ResponseStream<E> where E: ImapClient {
    pub fn new(future: Send<E::Transport>, state: ClientState,
               request_id: RequestId, next_state: Option<State>) -> Self {
        Self {
            future: Some(future),
            transport: None,
            state: Some(state),
            request_id: request_id,
            next_state: next_state,
            done: false,
        }
    }
}

impl<E> StateStream for ResponseStream<E> where E: ImapClient {
    type Item = ResponseData;
    type State = E;
    type Error = io::Error;
    fn poll(&mut self) -> Poll<StreamEvent<Self::Item, Self::State>, Self::Error> {
        match self.future.take() {
            Some(mut future) => match future.poll() {
                Ok(Async::Ready(transport)) => {
                    self.transport = Some(transport);
                },
                Ok(Async::NotReady) => {
                    self.future = Some(future);
                    return Ok(Async::NotReady);
                },
                Err(e) => {
                    return Err(e);
                },
            },
            None => {},
        }
        let mut transport = match self.transport.take() {
            None => return Ok(Async::NotReady),
            Some(mut transport) => transport,
        };
        if self.done {
            let mut state = self.state.take().unwrap(); // safe: initialized from start
            if let Some(next_state) = self.next_state.take() {
                state.state = next_state;
            }
            return Ok(Async::Ready(StreamEvent::Done(E::rebuild(transport, state))));
        }
        match transport.poll() {
            Ok(Async::Ready(Some(rsp))) => {
                if let Some(req_id) = rsp.request_id() {
                    self.done = *req_id == self.request_id;
                };
                self.transport = Some(transport);
                return Ok(Async::Ready(StreamEvent::Next(rsp)));
            },
            Err(e) => {
                return Err(e);
            },
            _ => (),
        }
        self.transport = Some(transport);
        Ok(Async::NotReady)
    }
}

pub enum ImapConnectFuture {
    #[doc(hidden)] TcpConnecting(ConnectFuture, String),
    #[doc(hidden)] TlsHandshake(ConnectAsync<TcpStream>),
    #[doc(hidden)] ServerGreeting(Option<ImapTls>),
}

impl Future for ImapConnectFuture {
    type Item = (Client, ResponseData);
    type Error = io::Error;
    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        let mut new = None;
        if let ImapConnectFuture::TcpConnecting(ref mut future, ref domain) = *self {
            let stream = try_ready!(future.poll());
            let ctx = TlsConnector::builder().unwrap().build().unwrap();
            let future = ctx.connect_async(domain, stream);
            new = Some(ImapConnectFuture::TlsHandshake(future));
        }
        if new.is_some() {
            *self = new.take().unwrap();
        }
        if let ImapConnectFuture::TlsHandshake(ref mut future) = *self {
            let transport = try_ready!(future.map_err(|e| {
                io::Error::new(io::ErrorKind::Other, e)
            }).poll()).framed(ImapCodec::default());
            new = Some(ImapConnectFuture::ServerGreeting(Some(transport)));
        }
        if new.is_some() {
            *self = new.take().unwrap();
        }
        if let ImapConnectFuture::ServerGreeting(ref mut wrapped) = *self {
            let msg = try_ready!(wrapped.as_mut().unwrap().poll()).unwrap();
            return Ok(Async::Ready((Client {
                transport: wrapped.take().unwrap(),
                state: ClientState::new(),
            }, msg)));
        }
        Ok(Async::NotReady)
    }
}

pub struct ClientState {
    state: State,
    request_ids: IdGenerator,
}

impl ClientState {
    pub fn new() -> Self {
        Self {
            state: State::NotAuthenticated,
            request_ids: IdGenerator::new(),
        }
    }
}

impl Default for ClientState {
    fn default() -> Self {
        Self::new()
    }
}

pub struct IdGenerator {
    next: u64,
}

impl IdGenerator {
    pub fn new() -> Self {
        Self { next: 0 }
    }
}

impl Default for IdGenerator {
    fn default() -> Self {
        Self::new()
    }
}

impl Iterator for IdGenerator {
    type Item = RequestId;
    fn next(&mut self) -> Option<Self::Item> {
        self.next += 1;
        Some(RequestId(format!("A{:04}", self.next % 10_000)))
    }
}
<|endoftext|>


<commit_before>use {QueryType, QueryClass, Name, Class, Header, RRData};

/// Parsed DNS packet
pub struct Packet<'a> {
    pub header: Header,
    pub questions: Vec<Question<'a>>,
    pub answers: Vec<ResourceRecord<'a>>,
    pub nameservers: Vec<ResourceRecord<'a>>,
    pub additional: Vec<ResourceRecord<'a>>,
}

/// A parsed chunk of data in the Query section of the packet
pub struct Question<'a> {
    pub qname: Name<'a>,
    pub qtype: QueryType,
    pub qclass: QueryClass,
}

/// A single DNS record
///
/// We aim to provide whole range of DNS records available. But as time is
/// limited we have some types of packets which are parsed and other provided
/// as unparsed slice of bytes.
pub struct ResourceRecord<'a> {
    pub name: Name<'a>,
    pub cls: Class,
    pub ttl: u32,
    pub data: RRData<'a>,
}
<commit_msg>Add `derive(Debug)` to structs<commit_after>use {QueryType, QueryClass, Name, Class, Header, RRData};

/// Parsed DNS packet
#[derive(Debug)]
pub struct Packet<'a> {
    pub header: Header,
    pub questions: Vec<Question<'a>>,
    pub answers: Vec<ResourceRecord<'a>>,
    pub nameservers: Vec<ResourceRecord<'a>>,
    pub additional: Vec<ResourceRecord<'a>>,
}

/// A parsed chunk of data in the Query section of the packet
#[derive(Debug)]
pub struct Question<'a> {
    pub qname: Name<'a>,
    pub qtype: QueryType,
    pub qclass: QueryClass,
}

/// A single DNS record
///
/// We aim to provide whole range of DNS records available. But as time is
/// limited we have some types of packets which are parsed and other provided
/// as unparsed slice of bytes.
#[derive(Debug)]
pub struct ResourceRecord<'a> {
    pub name: Name<'a>,
    pub cls: Class,
    pub ttl: u32,
    pub data: RRData<'a>,
}
<|endoftext|>


<commit_before>use std::rand::Rng;

use CliResult;
use config::{Config, Delimiter};
use util;

static USAGE: &'static str = "
Randomly samples CSV data uniformly using memory proportional to the size of
the sample.

This command is intended to provide a means to sample from a CSV data set that
is too big to fit into memory (for example, for use with commands like 'xsv
frequency' or 'xsv stats'). It will however visit every CSV record exactly
once, which is necessary to provide a uniform random sample. If you wish to
limit the number of records visited, use the 'xsv slice' command to pipe into
'xsv sample'.

Usage:
    xsv sample [options] <sample-size> [<input>]
    xsv sample --help

Common options:
    -h, --help             Display this message
    -o, --output <file>    Write output to <file> instead of stdout.
    -n, --no-headers       When set, the first row will be consider as part of
                           the population to sample from. (When not set, the
                           first row is the header row and will always appear
                           in the output.)
    -d, --delimiter <arg>  The field delimiter for reading CSV data.
                           Must be a single character. (default: ,)
";

#[deriving(Decodable)]
struct Args {
    arg_input: Option<String>,
    arg_sample_size: uint,
    flag_output: Option<String>,
    flag_no_headers: bool,
    flag_delimiter: Option<Delimiter>,
}

pub fn run(argv: &[&str]) -> CliResult<()> {
    let args: Args = try!(util::get_args(USAGE, argv));
    let rconfig = Config::new(&args.arg_input)
                         .delimiter(args.flag_delimiter)
                         .no_headers(args.flag_no_headers);

    let mut rdr = try!(rconfig.reader());
    let mut wtr = try!(Config::new(&args.flag_output).writer());

    // The following algorithm has been adapted from:
    // http://en.wikipedia.org/wiki/Reservoir_sampling

    // Fill the reservoir.
    let mut reservoir = Vec::with_capacity(args.arg_sample_size);
    {
        // Scope `rdr.byte_records()` so that the mutable borrow gets dropped.
        let mut records = rdr.byte_records().enumerate();
        for (_, row) in records.by_ref().take(reservoir.capacity()) {
            reservoir.push(try!(row));
        }

        // Now do the sampling.
        let mut rng = ::std::rand::task_rng();
        for (i, row) in records {
            let random = rng.gen_range(0, i+1);
            if random < args.arg_sample_size {
                reservoir[random] = try!(row);
            }
        }
    }

    try!(rconfig.write_headers(&mut rdr, &mut wtr));
    for row in reservoir.into_iter() {
        try!(wtr.write_bytes(row.into_iter()));
    }
    Ok(try!(wtr.flush()))
}
<commit_msg>Add index optimization to 'sample' command.<commit_after>use std::rand::Rng;

use csv::{mod, ByteString};
use csv::index::Indexed;

use CliResult;
use config::{Config, Delimiter};
use util;

static USAGE: &'static str = "
Randomly samples CSV data uniformly using memory proportional to the size of
the sample.

When an index is present, this command will use random indexing if the sample
size is less than 10% of the total number of records. This allows for efficient
sampling such that the entire CSV file is not parsed.

This command is intended to provide a means to sample from a CSV data set that
is too big to fit into memory (for example, for use with commands like 'xsv
frequency' or 'xsv stats'). It will however visit every CSV record exactly
once, which is necessary to provide a uniform random sample. If you wish to
limit the number of records visited, use the 'xsv slice' command to pipe into
'xsv sample'.

Usage:
    xsv sample [options] <sample-size> [<input>]
    xsv sample --help

Common options:
    -h, --help             Display this message
    -o, --output <file>    Write output to <file> instead of stdout.
    -n, --no-headers       When set, the first row will be consider as part of
                           the population to sample from. (When not set, the
                           first row is the header row and will always appear
                           in the output.)
    -d, --delimiter <arg>  The field delimiter for reading CSV data.
                           Must be a single character. (default: ,)
";

#[deriving(Decodable)]
struct Args {
    arg_input: Option<String>,
    arg_sample_size: u64,
    flag_output: Option<String>,
    flag_no_headers: bool,
    flag_delimiter: Option<Delimiter>,
}

pub fn run(argv: &[&str]) -> CliResult<()> {
    let args: Args = try!(util::get_args(USAGE, argv));
    let rconfig = Config::new(&args.arg_input)
                         .delimiter(args.flag_delimiter)
                         .no_headers(args.flag_no_headers);
    let sample_size = args.arg_sample_size;

    let mut wtr = try!(Config::new(&args.flag_output).writer());
    let sampled = match try!(rconfig.indexed()) {
        Some(mut idx) => {
            if do_random_access(sample_size, idx.count()) {
                try!(rconfig.write_headers(idx.csv(), &mut wtr));
                try!(sample_random_access(&mut idx, sample_size))
            } else {
                let mut rdr = try!(rconfig.reader());
                try!(rconfig.write_headers(&mut rdr, &mut wtr));
                try!(sample_reservoir(&mut rdr, sample_size))
            }
        }
        _ => {
            let mut rdr = try!(rconfig.reader());
            try!(rconfig.write_headers(&mut rdr, &mut wtr));
            try!(sample_reservoir(&mut rdr, sample_size))
        }
    };
    for row in sampled.into_iter() {
        try!(wtr.write_bytes(row.into_iter()));
    }
    Ok(try!(wtr.flush()))
}

fn sample_random_access<R: Reader + Seek, I: Reader + Seek>
                       (idx: &mut Indexed<R, I>, sample_size: u64)
                       -> CliResult<Vec<Vec<ByteString>>> {
    let mut all_indices = range(0, idx.count()).collect::<Vec<_>>();
    let mut rng = ::std::rand::task_rng();
    rng.shuffle(all_indices.as_mut_slice());

    let mut sampled = Vec::with_capacity(sample_size as uint);
    for i in all_indices.into_iter().take(sample_size as uint) {
        try!(idx.seek(i));
        let mut rdr = idx.csv();
        sampled.push(try!(rdr.byte_records().next().unwrap()));
    }
    Ok(sampled)
}

fn sample_reservoir<R: Reader>
                   (rdr: &mut csv::Reader<R>, sample_size: u64)
                   -> CliResult<Vec<Vec<ByteString>>> {
    // The following algorithm has been adapted from:
    // http://en.wikipedia.org/wiki/Reservoir_sampling
    let mut reservoir = Vec::with_capacity(sample_size as uint);
    let mut records = rdr.byte_records().enumerate();
    for (_, row) in records.by_ref().take(reservoir.capacity()) {
        reservoir.push(try!(row));
    }

    // Now do the sampling.
    let mut rng = ::std::rand::task_rng();
    for (i, row) in records {
        let random = rng.gen_range(0, i+1);
        if random < sample_size as uint {
            reservoir[random] = try!(row);
        }
    }
    Ok(reservoir)
}

fn do_random_access(sample_size: u64, total: u64) -> bool {
    sample_size <= (total / 10)
}
<|endoftext|>