add functions and variables to language

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)]
-pub enum ParseTree {
-    Leaf(Token),
-    Branch(Op, Box<ParseTree>, Box<ParseTree>),
+#[derive(Clone, Debug)]
+pub enum ParseTree<'a> {
+    Add(Box<ParseTree<'a>>, Box<ParseTree<'a>>),
+    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 {
-    pub fn new<'a, I: Iterator<Item = &'a Token>>(tokens: &mut I) -> Result<Self, ParseError> {
+impl<'a> ParseTree<'a> {
+    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())),
-                Token::Operator(op) => {
-                    let left = ParseTree::new(tokens)?;
-                    let right = ParseTree::new(tokens)?;
+                // Just return scalars
+                Token::Scalar(x) => Ok(ParseTree::Scalar(*x)),
 
-                    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::Leaf(Token::Identifier(_)) => unimplemented!(),
-            ParseTree::Leaf(Token::Operator(_)) => panic!("This absolutely should not happen"),
-            ParseTree::Branch(op, left, right) => match op {
-                Op::Add => left.evaluate() + right.evaluate(),
-                Op::Sub => left.evaluate() - right.evaluate(),
-                Op::Mul => left.evaluate() * right.evaluate(),
-                Op::Div => left.evaluate() / right.evaluate(),
-                Op::Exp => left.evaluate().powf(right.evaluate()),
+            ParseTree::Add(l, r) => Ok(l.evaluate(globals, locals)? + r.evaluate(globals, locals)?),
+            ParseTree::Sub(l, r) => Ok(l.evaluate(globals, locals)? + r.evaluate(globals, locals)?),
+            ParseTree::Mul(l, r) => Ok(l.evaluate(globals, locals)? * r.evaluate(globals, locals)?),
+            ParseTree::Div(l, r) => Ok(l.evaluate(globals, locals)? / r.evaluate(globals, locals)?),
+            ParseTree::Exp(l, r)
+                => Ok(l.evaluate(globals, locals)?.powf(r.evaluate(globals, locals)?)),
+            ParseTree::Equ(ident, value, body) => {
+                let value = value.evaluate(globals, locals)?;
+
+                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"));
+    }
 }