add functions and variables to language

src/commands/eval/mod.rs

@@ -1,5 +1,8 @@
 use crate::common::{Context, Error};
 
+use std::collections::HashMap;
+use std::borrow::Cow;
+
 mod tokenizer;
 mod parse;
 
@@ -7,9 +10,17 @@ fn evaluate(expr: &str) -> Result<f64, Error> {
     let tokens = tokenizer::Token::tokenize(expr)?;
     let mut tokens = tokens.iter();
 
-    let tree = parse::ParseTree::new(&mut tokens)?;
+    let globals = HashMap::new();
+    let locals = HashMap::new();
+    let mut locals = Cow::Borrowed(&locals);
 
-    Ok(tree.evaluate())
+    let tree = parse::ParseTree::parse(&mut tokens, &globals, &mut locals)?;
+
+    let mut globals = HashMap::new();
+    let locals = HashMap::new();
+    let mut locals = Cow::Borrowed(&locals);
+
+    tree.evaluate(&mut globals, &mut locals)
 }
 
 /// Evaluates an expression (uses Polish Notation)

src/commands/eval/parse.rs

@@ -1,38 +1,192 @@
 use std::error::Error;
+use std::borrow::Cow;
+use std::collections::HashMap;
 use std::fmt::Display;
-use super::tokenizer::{Token, Op};
+use crate::common;
+use super::tokenizer::{self, Token};
 
-#[derive(Debug)]
+#[derive(Clone, Debug)]
-pub enum ParseTree {
+pub enum ParseTree<'a> {
-    Leaf(Token),
+    Add(Box<ParseTree<'a>>, Box<ParseTree<'a>>),
-    Branch(Op, Box<ParseTree>, Box<ParseTree>),
+    Sub(Box<ParseTree<'a>>, Box<ParseTree<'a>>),
+    Mul(Box<ParseTree<'a>>, Box<ParseTree<'a>>),
+    Div(Box<ParseTree<'a>>, Box<ParseTree<'a>>),
+    Exp(Box<ParseTree<'a>>, Box<ParseTree<'a>>),
+    Equ(&'a str, Box<ParseTree<'a>>, Box<ParseTree<'a>>),
+    GlobalEqu(&'a str, Box<ParseTree<'a>>),
+    Compose(Box<ParseTree<'a>>, Box<ParseTree<'a>>),
+    Id(Box<ParseTree<'a>>),
+    FunctionDeclaration(&'a str, Vec<Object<'a>>, Box<ParseTree<'a>>, Box<ParseTree<'a>>),
+    FunctionApplication(&'a str, Vec<ParseTree<'a>>),
+    Variable(&'a str),
+    Scalar(f64),
+}
+
+#[derive(Clone, Debug)]
+pub struct FunctionDeclaration<'a> {
+    name: &'a str,
+    args: Vec<Object<'a>>,
+}
+
+#[derive(Clone, Debug)]
+pub struct Function<'a> {
+    decl: FunctionDeclaration<'a>,
+    body: Option<Box<ParseTree<'a>>>, // may be used in declarations where a value isn't provided
+}
+
+#[derive(Clone, Debug)]
+pub struct Variable<'a> {
+    name: &'a str,
+    body: Option<Box<ParseTree<'a>>>, // may be used in declarations where a value isn't provided
+}
+
+impl<'a> Variable<'a> {
+    pub fn new(name: &'a str, body: Option<Box<ParseTree<'a>>>) -> Self {
+        Self { name, body }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub enum Object<'a> {
+    Variable(Variable<'a>),
+    Func(Function<'a>),
 }
 
 #[derive(Debug, Clone)]
 pub enum ParseError {
     UnexpectedEndInput,
+    IdentifierUndefined(String),
+    InvalidIdentifier,
+    FunctionUndefined,
+    VariableUndefined,
 }
 
 impl Display for ParseError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             ParseError::UnexpectedEndInput => write!(f, "Input ended unexpectedly"),
+            ParseError::IdentifierUndefined(name) => write!(f, "Undefined variable `{}`", name),
+            ParseError::InvalidIdentifier => write!(f, "Invalid identifier"),
+            ParseError::FunctionUndefined => write!(f, "Undefined function"),
+            ParseError::VariableUndefined => write!(f, "Undefined variable"),
         }
     }
 }
 
 impl Error for ParseError {}
 
-impl ParseTree {
+impl<'a> ParseTree<'a> {
-    pub fn new<'a, I: Iterator<Item = &'a Token>>(tokens: &mut I) -> Result<Self, ParseError> {
+    pub fn parse<I: Iterator<Item = &'a Token>>(
+        tokens: &mut I,
+        globals: &HashMap<String, FunctionDeclaration<'a>>,
+        locals: &mut Cow<HashMap<String, FunctionDeclaration<'a>>>) -> Result<Self, ParseError>
+    {
         if let Some(token) = tokens.next() {
             match token {
-                Token::Scalar(_) | Token::Identifier(_) => Ok(Self::Leaf(token.clone())),
+                // Just return scalars
-                Token::Operator(op) => {
+                Token::Scalar(x) => Ok(ParseTree::Scalar(*x)),
-                    let left = ParseTree::new(tokens)?;
-                    let right = ParseTree::new(tokens)?;
 
-                    Ok(Self::Branch(op.clone(), Box::new(left), Box::new(right)))
+                // Get any identifiers as objects from local first then global scope
+                Token::Identifier(ident) => {
+                    // If it is found to be a function, get its argument count.
+                    // During parsing, we only keep track of function definitions
+                    // so that we know how many arguments it takes
+                    if let Some(decl) = locals.clone().get(ident).or(globals.get(ident)) {
+                        let args = decl.args.iter()
+                            .map(|_| ParseTree::parse(tokens, globals, locals)).collect::<Result<Vec<_>, ParseError>>()?;
+
+                        Ok(ParseTree::FunctionApplication(ident, args))
+                    } else {
+                        Ok(ParseTree::Variable(ident))
+                    }
+                }
+
+                Token::Operator(op) => {
+                    match op {
+                        tokenizer::Op::Add => Ok(ParseTree::Add(
+                            Box::new(ParseTree::parse(tokens, globals, locals)?),
+                            Box::new(ParseTree::parse(tokens, globals, locals)?)
+                        )),
+                        tokenizer::Op::Sub => Ok(ParseTree::Sub(
+                            Box::new(ParseTree::parse(tokens, globals, locals)?),
+                            Box::new(ParseTree::parse(tokens, globals, locals)?)
+                        )),
+                        tokenizer::Op::Mul => Ok(ParseTree::Mul(
+                            Box::new(ParseTree::parse(tokens, globals, locals)?),
+                            Box::new(ParseTree::parse(tokens, globals, locals)?)
+                        )),
+                        tokenizer::Op::Div => Ok(ParseTree::Div(
+                            Box::new(ParseTree::parse(tokens, globals, locals)?),
+                            Box::new(ParseTree::parse(tokens, globals, locals)?)
+                        )),
+                        tokenizer::Op::Exp => Ok(ParseTree::Exp(
+                            Box::new(ParseTree::parse(tokens, globals, locals)?),
+                            Box::new(ParseTree::parse(tokens, globals, locals)?)
+                        )),
+                        tokenizer::Op::Equ | tokenizer::Op::LazyEqu => {
+                            let token = tokens.next().ok_or(ParseError::UnexpectedEndInput)?;
+
+                            if let Token::Identifier(ident) = token {
+                                Ok(ParseTree::Equ(
+                                    ident,
+                                    Box::new(ParseTree::parse(tokens, globals, locals)?),
+                                    Box::new(ParseTree::parse(tokens, globals, locals)?)))
+                            } else {
+                                Err(ParseError::InvalidIdentifier)
+                            }
+                        }
+                        tokenizer::Op::GlobalEqu | tokenizer::Op::LazyGlobalEqu => {
+                            let token = tokens.next().ok_or(ParseError::UnexpectedEndInput)?;
+
+                            if let Token::Identifier(ident) = token {
+                                Ok(ParseTree::GlobalEqu(
+                                    ident,
+                                    Box::new(ParseTree::parse(tokens, globals, locals)?)
+                                ))
+                            } else {
+                                Err(ParseError::InvalidIdentifier)
+                            }
+                        }
+                        tokenizer::Op::FunctionDeclare(arg_count) => {
+                            let token = tokens.next().ok_or(ParseError::UnexpectedEndInput)?;
+                            
+                            if let Token::Identifier(ident) = token {
+                                let args: Vec<Object> = tokens.take(*arg_count)
+                                    .map(|token| match token {
+                                        Token::Identifier(s)
+                                            => Ok(Object::Variable(Variable::new(s, None))),
+                                        _ => Err(ParseError::InvalidIdentifier),
+                                    }).collect::<Result<_, ParseError>>()?;
+
+                                if args.len() < *arg_count {
+                                    return Err(ParseError::InvalidIdentifier);
+                                }
+
+                                let locals = locals.to_mut();
+
+                                locals.insert(ident.clone(), FunctionDeclaration {
+                                    name: ident,
+                                    args: args.clone()
+                                });
+
+                                Ok(ParseTree::FunctionDeclaration(
+                                    ident,
+                                    args,
+                                    Box::new(ParseTree::parse(tokens, globals, &mut Cow::Borrowed(&*locals))?),
+                                    Box::new(ParseTree::parse(tokens, globals, &mut Cow::Borrowed(&*locals))?)))
+                            } else {
+                                Err(ParseError::InvalidIdentifier)
+                            }
+                        }
+                        tokenizer::Op::Compose => {
+                            Ok(ParseTree::Compose(
+                                Box::new(ParseTree::parse(tokens, globals, locals)?),
+                                Box::new(ParseTree::parse(tokens, globals, locals)?)
+                            ))
+                        }
+                        tokenizer::Op::Id =>
+                            Ok(ParseTree::Id(Box::new(ParseTree::parse(tokens, globals, locals)?)))
+                    }
                 }
             }
         } else {
@@ -40,18 +194,121 @@ impl ParseTree {
         }
     }
 
-    pub fn evaluate(&self) -> f64 {
+    pub fn evaluate(
+        self,
+        globals: &mut HashMap<String, Object<'a>>,
+        locals: &mut Cow<HashMap<String, Object<'a>>>) -> Result<f64, common::Error>
+    {
         match self {
-            ParseTree::Leaf(Token::Scalar(value)) => *value,
+            ParseTree::Add(l, r) => Ok(l.evaluate(globals, locals)? + r.evaluate(globals, locals)?),
-            ParseTree::Leaf(Token::Identifier(_)) => unimplemented!(),
+            ParseTree::Sub(l, r) => Ok(l.evaluate(globals, locals)? + r.evaluate(globals, locals)?),
-            ParseTree::Leaf(Token::Operator(_)) => panic!("This absolutely should not happen"),
+            ParseTree::Mul(l, r) => Ok(l.evaluate(globals, locals)? * r.evaluate(globals, locals)?),
-            ParseTree::Branch(op, left, right) => match op {
+            ParseTree::Div(l, r) => Ok(l.evaluate(globals, locals)? / r.evaluate(globals, locals)?),
-                Op::Add => left.evaluate() + right.evaluate(),
+            ParseTree::Exp(l, r)
-                Op::Sub => left.evaluate() - right.evaluate(),
+                => Ok(l.evaluate(globals, locals)?.powf(r.evaluate(globals, locals)?)),
-                Op::Mul => left.evaluate() * right.evaluate(),
+            ParseTree::Equ(ident, value, body) => {
-                Op::Div => left.evaluate() / right.evaluate(),
+                let value = value.evaluate(globals, locals)?;
-                Op::Exp => left.evaluate().powf(right.evaluate()),
+
+                let locals = locals.to_mut();
+
+                locals.insert(ident.to_string(),
+                              Object::Variable(
+                                  Variable::new(ident, Some(Box::new(ParseTree::Scalar(value))))));
+
+                body.evaluate(globals, &mut Cow::Borrowed(&locals))
+            }
+            ParseTree::GlobalEqu(ident, body) => {
+                globals.insert(ident.to_string(),
+                               Object::Variable(Variable::new(ident, Some(body.clone()))));
+
+                Ok(0.0)
+            }
+            ParseTree::Compose(l, r) => {
+                let _ = l.evaluate(globals, locals);
+                r.evaluate(globals, locals)
+            }
+            ParseTree::Id(body) => body.evaluate(globals, locals),
+            ParseTree::FunctionDeclaration(name, args, body, cont) => {
+                let locals = locals.to_mut();
+
+                locals.insert(name.to_string(), Object::Func(Function {
+                    decl: FunctionDeclaration {
+                        name,
+                        args: args.clone(),
+                    },
+                    body: Some(body.clone())
+                }));
+
+                cont.evaluate(globals, &mut Cow::Borrowed(&locals))
+            }
+            ParseTree::FunctionApplication(name, params) => {
+                let locals = locals.to_mut();
+                let obj = locals.get(name).or(globals.get(name)).cloned();
+
+                if let Some(Object::Func(func)) = obj {
+                    for (param, arg) in params.iter().zip(func.decl.args.iter()) {
+                        match arg {
+                            Object::Variable(v)
+                                => locals.insert(
+                                    v.name.to_string(),
+                                    Object::Variable(
+                                        Variable::new(
+                                            &v.name,
+                                            Some(Box::new(
+                                                ParseTree::Scalar(
+                                                    param.clone().evaluate(
+                                                        globals, &mut Cow::Borrowed(&locals))?)))))),
+                            Object::Func(func)
+                                => locals.insert(
+                                    func.decl.name.to_string(),
+                                    Object::Func(func.clone()))
+                        };
+                    }
+
+                    func.body.ok_or(ParseError::FunctionUndefined.into())
+                        .and_then(|body|
+                            body.clone().evaluate(globals, &mut Cow::Borrowed(&locals)))
+                } else {
+                    Err(ParseError::FunctionUndefined.into())
+                }
+            }
+            ParseTree::Variable(ident) => {
+                let locals = locals.to_mut();
+                let obj = locals.get(ident).or(globals.get(ident)).cloned();
+
+                if let Some(Object::Variable(obj)) = obj {
+                    return obj.body.clone().ok_or(ParseError::VariableUndefined.into())
+                        .and_then(|body| body.clone().evaluate(globals, &mut Cow::Borrowed(&locals)));
+                }
+
+                Err(ParseError::IdentifierUndefined(ident.to_string()).into())
+            }
+            ParseTree::Scalar(x) => Ok(x),
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_parse() {
+        let tokens = Token::tokenize("= x 15 : square x ** x 2 square x").expect("failed to tokenize");
+        let mut tokens = tokens.iter();
+
+        let globals = HashMap::new();
+        let locals = HashMap::new();
+        let mut locals = Cow::Borrowed(&locals);
+
+        let tree = ParseTree::parse(&mut tokens, &globals, &mut locals).expect("failed to parse");
+
+        eprintln!("{tree:?}");
+
+        let mut globals = HashMap::new();
+        let locals = HashMap::new();
+        let mut locals = Cow::Borrowed(&locals);
+
+        eprintln!("{}", tree.evaluate(&mut globals, &mut locals).expect("failed to evaluate"));
+    }
 }

src/commands/eval/tokenizer.rs

@@ -1,4 +1,5 @@
 use std::error;
+
 use crate::common::Error;
 use std::fmt::{Display, Formatter};
 use std::str::FromStr;
@@ -10,6 +11,13 @@ pub enum Op {
     Mul,
     Div,
     Exp,
+    Equ,
+    LazyEqu,
+    GlobalEqu,
+    LazyGlobalEqu,
+    FunctionDeclare(usize),
+    Compose,
+    Id,
 }
 
 #[derive(Debug, Clone)]
@@ -43,44 +51,46 @@ impl Token {
         Token::Scalar(value)
     }
 
-    pub fn add() -> Token {
+    pub fn operator(op: Op) -> Token {
-        Token::Operator(Op::Add)
+        Token::Operator(op)
     }
-
-    pub fn sub() -> Token {
-        Token::Operator(Op::Sub)
-    }
-
-    pub fn mul() -> Token {
-        Token::Operator(Op::Mul)
-    }
-
-    pub fn div() -> Token {
-        Token::Operator(Op::Div)
-    }
-
-    pub fn exp() -> Token {
-        Token::Operator(Op::Exp)
-    }
-
     pub fn tokenize(s: &str) -> Result<Vec<Self>, Error> {
         s.split_whitespace().map(Token::from_str).collect()
     }
 }
 
+fn get_dot_count<I: Iterator<Item = char>>(s: I) -> usize {
+    s.fold(0, |acc, c| acc + match c {
+        ':' => 2,
+        '.' => 1,
+        _ => 0,
+    })
+}
+
 impl FromStr for Token {
     type Err = Error;
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
-        match s {
+        let s = if s.starts_with("\\") { s.chars().skip(1).collect() } else { s.to_string() };
+
+        match s.as_str() {
             // First check if s is an operator
-            "+"  => Ok(Token::add()),
+            "+"  => Ok(Token::operator(Op::Add)),
-            "-"  => Ok(Token::sub()),
+            "-"  => Ok(Token::operator(Op::Sub)),
-            "*"  => Ok(Token::mul()),
+            "*"  => Ok(Token::operator(Op::Mul)),
-            "**" => Ok(Token::exp()),
+            "**" => Ok(Token::operator(Op::Exp)),
-            "/"  => Ok(Token::div()),
+            "/"  => Ok(Token::operator(Op::Div)),
+            "="  => Ok(Token::operator(Op::Equ)),
+            "."  => Ok(Token::operator(Op::LazyEqu)),
+            "=>" => Ok(Token::operator(Op::GlobalEqu)),
+            ".>" => Ok(Token::operator(Op::LazyGlobalEqu)),
+            "~"  => Ok(Token::operator(Op::Compose)),
+            "," => Ok(Token::operator(Op::Id)),
             _ => {
-                if s.starts_with(|c| char::is_digit(c, 10)) {
+                // variable length operators
+                if s.starts_with(':') {
+                    Ok(Token::operator(Op::FunctionDeclare(1 + get_dot_count(s[1..].chars()))))
+                } else if s.starts_with(|c| char::is_digit(c, 10)) {
                     Ok(Token::scalar(s.parse()?))
                 } else if s.starts_with(char::is_alphabetic)
                         && s.chars().skip(1).all(char::is_alphanumeric) {