Implement Optionally-Typed Function Parameters (#252)

examples/fn.ion

@@ -13,3 +13,11 @@ fn another_test
 end
 
 another_test
+
+fn square n:int
+    calc $n \* $n
+end
+
+for num in 1 2 3 4 5 a 1.5
+    square $num
+end

examples/fn.out

@@ -10,3 +10,8 @@ goodbye
 7
 8
 9
+1
+4
+9
+16
+25

src/parser/grammar.rustpeg

@@ -1,5 +1,6 @@
 use parser::assignments::parse_assignment;
 use parser::{pipelines, ArgumentSplitter};
+use parser::peg::get_function_args;
 use shell::flow_control::{ElseIf, Statement};
 
 #[pub]
@@ -72,13 +73,13 @@ end_ -> Statement
 
 #[pub]
 fn_ -> Statement
-    = whitespace* "fn " n:_name whitespace* args:_args whitespace* description:_description? {
-        Statement::Function {
+    = whitespace* "fn " n:_name whitespace* args:_fn_args whitespace* description:_description? {?
+        get_function_args(args).map(|args| Statement::Function {
             description: description.unwrap_or("".into()),
             name: n.into(),
             args: args,
             statements: Vec::new(),
-        }
+        }).ok_or("ion: invalid function argument\n")
     }
 
 _description -> String
@@ -87,11 +88,17 @@ _description -> String
 _name -> String
       = n:$([A-z0-9_]+) { n.into() }
 
+_fn_args -> Vec<String>
+    = _fn_arg ** " "
+
 _args -> Vec<String>
       = _arg ** " "
 
 _arg -> String
-     = n:$([A-z0-9]+) { n.into() }
+     = n:$([A-z0-9_]+) { n.into() }
+
+_fn_arg -> String
+    = n:$([A-z0-9_:]+) { n.into()}
 
 #[pub]
 for_ -> Statement

src/parser/peg.rs

@@ -1,6 +1,6 @@
 use std::io::{stderr, Write};
 
-use shell::flow_control::Statement;
+use shell::flow_control::{Statement, FunctionArgument, Type};
 use self::grammar::parse_;
 use shell::directory_stack::DirectoryStack;
 use shell::Job;
@@ -61,6 +61,43 @@ pub fn parse(code: &str) -> Statement {
 	}
 }
 
+pub fn get_function_args(args: Vec<String>) -> Option<Vec<FunctionArgument>> {
+    let mut fn_args = Vec::with_capacity(args.len());
+    for argument in args.into_iter() {
+        let length = argument.len();
+        let argument = if argument.ends_with(":int") {
+            if length <= 4 { return None }
+            let arg = &argument[..length-4];
+            if arg.contains(':') { return None }
+            FunctionArgument::Typed (
+                arg.to_owned(),
+                Type::Int
+            )
+        } else if argument.ends_with(":float") {
+            if length <= 6 { return None }
+            let arg = &argument[..length-6];
+            if arg.contains(':') { return None }
+            FunctionArgument::Typed (
+                arg.to_owned(),
+                Type::Float
+            )
+        } else if argument.ends_with(":bool") {
+            if length <= 5 { return None }
+            let arg = &argument[..length-5];
+            if arg.contains(':') { return None }
+            FunctionArgument::Typed (
+                arg.to_owned(),
+                Type::Bool
+            )
+        } else {
+            FunctionArgument::Untyped(argument)
+        };
+        fn_args.push(argument);
+    }
+
+    Some(fn_args)
+}
+
 mod grammar {
     include!(concat!(env!("OUT_DIR"), "/grammar.rs"));
 }

src/shell/flow_control.rs

@@ -8,6 +8,13 @@ pub struct ElseIf {
     pub success:    Vec<Statement>
 }
 
+#[derive(Debug, PartialEq, Clone, Copy)]
+pub enum Type { Float, Int, Bool }
+
+#[derive(Debug, PartialEq, Clone)]
+pub enum FunctionArgument { Typed(String, Type), Untyped(String) }
+
+
 #[derive(Debug, PartialEq, Clone)]
 pub enum Statement {
     Let {
@@ -23,7 +30,7 @@ pub enum Statement {
     Function {
         name: Identifier,
         description: String,
-        args: Vec<String>,
+        args: Vec<FunctionArgument>,
         statements: Vec<Statement>
     },
     For {
@@ -47,7 +54,7 @@ pub enum Statement {
 pub struct FlowControl {
     pub level:             usize,
     pub current_statement: Statement,
-    pub current_if_mode:      u8 // { 0 = SUCCESS; 1 = FAILURE }
+    pub current_if_mode:   u8 // { 0 = SUCCESS; 1 = FAILURE }
 }
 
 impl Default for FlowControl {
@@ -64,13 +71,15 @@ impl Default for FlowControl {
 pub struct Function {
     pub description: String,
     pub name: Identifier,
-    pub args: Vec<String>,
+    pub args: Vec<FunctionArgument>,
     pub statements: Vec<Statement>
 }
 
-pub fn collect_loops<I>(iterator: &mut I, statements: &mut Vec<Statement>, level: &mut usize)
-    where I: Iterator<Item = Statement>
-{
+pub fn collect_loops <I: Iterator<Item = Statement>> (
+    iterator: &mut I,
+    statements: &mut Vec<Statement>,
+    level: &mut usize
+) {
     #[allow(while_let_on_iterator)]
     while let Some(statement) = iterator.next() {
         match statement {

src/shell/mod.rs

@@ -31,7 +31,7 @@ use types::*;
 use smallstring::SmallString;
 use self::completer::{MultiCompleter, IonFileCompleter};
 use self::directory_stack::DirectoryStack;
-use self::flow_control::{FlowControl, Function};
+use self::flow_control::{FlowControl, Function, FunctionArgument, Type};
 use self::variables::Variables;
 use self::status::*;
 use self::pipe::execute_pipeline;
@@ -324,8 +324,6 @@ impl<'a> Shell<'a> {
         }
     }
 
-
-
     /// Executes a pipeline and returns the final exit status of the pipeline.
     /// To avoid infinite recursion when using aliases, the noalias boolean will be set the true
     /// if an alias branch was executed.
@@ -370,33 +368,72 @@ impl<'a> Shell<'a> {
                 if pipeline.jobs[0].args.len() - 1 == function.args.len() {
                     let mut variables_backup: FnvHashMap<&str, Option<Value>> =
                         FnvHashMap::with_capacity_and_hasher (
-                        64, Default::default()
-                    );
-                    for (name, value) in function.args.iter().zip(pipeline.jobs[0].args.iter().skip(1)) {
+                            64, Default::default()
+                        );
+
+                    let mut bad_argument: Option<(&str, Type)> = None;
+                    for (name_arg, value) in function.args.iter().zip(pipeline.jobs[0].args.iter().skip(1)) {
+                        let name: &str = match name_arg {
+                            &FunctionArgument::Typed(ref name, ref type_) => {
+                                match *type_ {
+                                    Type::Float if value.parse::<f64>().is_ok() => name.as_str(),
+                                    Type::Int if value.parse::<i64>().is_ok() => name.as_str(),
+                                    Type::Bool if value == "true" || value == "false" => name.as_str(),
+                                    _ => {
+                                        bad_argument = Some((value.as_str(), *type_));
+                                        break
+                                    }
+                                }
+                            },
+                            &FunctionArgument::Untyped(ref name) => name.as_str()
+                        };
                         variables_backup.insert(name, self.variables.get_var(name));
                         self.variables.set_var(name, value);
                     }
 
-                    self.execute_statements(function.statements);
+                    match bad_argument {
+                        Some((actual_value, expected_type)) => {
+                            for (name, value_option) in &variables_backup {
+                                match *value_option {
+                                    Some(ref value) => self.variables.set_var(name, value),
+                                    None => {self.variables.unset_var(name);},
+                                }
+                            }
+
+                            let type_ = match expected_type {
+                                Type::Float => "Float",
+                                Type::Int   => "Int",
+                                Type::Bool  => "Bool"
+                            };
+
+                            let stderr = io::stderr();
+                            let mut stderr = stderr.lock();
+                            let _ = writeln!(stderr, "ion: function argument has invalid type: expected {}, found value \'{}\'", type_, actual_value);
+                            Some(FAILURE)
+                        }
+                        None => {
+                            self.execute_statements(function.statements);
 
-                    for (name, value_option) in &variables_backup {
-                        match *value_option {
-                            Some(ref value) => self.variables.set_var(name, value),
-                            None => {self.variables.unset_var(name);},
+                            for (name, value_option) in &variables_backup {
+                                match *value_option {
+                                    Some(ref value) => self.variables.set_var(name, value),
+                                    None => {self.variables.unset_var(name);},
+                                }
+                            }
+                            None
                         }
                     }
-                    None
                 } else {
                     let stderr = io::stderr();
                     let mut stderr = stderr.lock();
-                    let _ = writeln!(stderr, "This function takes {} arguments, but you provided {}",
+                    let _ = writeln!(stderr, "ion: function takes {} arguments, but you provided {}",
                         function.args.len(), pipeline.jobs[0].args.len()-1);
                     Some(NO_SUCH_COMMAND) // not sure if this is the right error code
                 }
             } else {
                 let stderr = io::stderr();
                 let mut stderr = stderr.lock();
-                let _ = writeln!(stderr, "Function pipelining is not implemented yet");
+                let _ = writeln!(stderr, "ion: function pipelining is not implemented yet");
                 Some(FAILURE)
             }
         // If not a shell command or a shell function execute the pipeline and set the exit_status