computorv1/src/parser/ast_builder.rs

use super::Token;

#[derive(Debug, PartialEq, Clone)]
pub struct Rational {
    numerator: i128,
    denominator: i128,
}

#[derive(Debug, PartialEq, Clone)]
pub struct GaussianRational {
    real: Rational,
    imaginary: Rational,
}


#[derive(Debug, PartialEq, Clone)]
pub enum Node {
    Leaf(Vec<GaussianRational>),
    Internal {
        operator: Token,
        lhs: Box<Node>,
        rhs: Box<Node>,
    }
}

/*
   pub struct Node {
   _lhs: Option<Box<Node>>,
   _rhs: Option<Box<Node>>,
   }
   */

fn zero() -> GaussianRational {
    GaussianRational {
        real: Rational {
            numerator: 0,
            denominator: 1,
        },
        imaginary: Rational {
            numerator: 0,
            denominator: 1,
        }
    }
}

fn minus_one() -> GaussianRational {
    GaussianRational {
        real: Rational {
            numerator: -1,
            denominator: 1,
        },
        imaginary: Rational {
            numerator: 0,
            denominator: 1,
        }
    }
}

fn one() -> GaussianRational {
    GaussianRational {
        real: Rational {
            numerator: 1,
            denominator: 1,
        },
        imaginary: Rational {
            numerator: 0,
            denominator: 1,
        }
    }
}

fn i() -> GaussianRational {
    GaussianRational {
        real: Rational {
            numerator: 0,
            denominator: 1,
        },
        imaginary: Rational {
            numerator: 1,
            denominator: 1,
        }
    }
}

fn minus_i() -> GaussianRational {
    GaussianRational {
        real: Rational {
            numerator: 0,
            denominator: 1,
        },
        imaginary: Rational {
            numerator: -1,
            denominator: 1,
        }
    }
}

fn x() -> Vec<GaussianRational> {
    vec![zero(), one()]
}

fn minus_x() -> Vec<GaussianRational> {
    vec![zero(), minus_one()]
}

fn insert_operator_node(mut tokens: Vec<Token>, position: usize) -> Node {
    let tokens_right = tokens.split_off(position + 1);
    let token = tokens.pop().unwrap();
    Node::Internal {
        operator: token,
        lhs: Box::new(build_ast(tokens)),
        rhs: Box::new(build_ast(tokens_right)),
    }
}

fn get_parentheses_map(tokens: &Vec<Token>) -> Vec<usize> {
    let mut parentheses_map = vec![];
    let mut score = 0;
    for token in tokens {
        match token {
            Token::OpenParenthesis() => score += 1,
            Token::CloseParenthesis() => score -= 1,
            _ => (),
        }
        parentheses_map.push(score);
    }
    parentheses_map
}

fn parse_number(string: &String, sign: i128) -> GaussianRational {
    let mut number = Rational {
        numerator: 0,
        denominator: 1,
    };
    let mut is_floating = false;

    for c in string.chars() {
        match c {
            '.' => is_floating = true,
            '0'..='9' => {
                number.numerator = 10 * number.numerator + String::from(c).parse::<i128>().unwrap();
                if is_floating == true {
                    number.denominator *= 10;
                }
            }
            _ => unreachable!(),
        }
    }

    number.numerator *= sign;

    GaussianRational {
        real: number,
        imaginary: Rational {
            numerator: 0,
            denominator: 1,
        }
    }
}

fn check_parentheses(tokens: &Vec<Token>) -> bool {
    let mut level = 0;
    let mut count = 0;
    for token in tokens {
        match token {
            Token::OpenParenthesis() => level += 1,
            Token::CloseParenthesis() => level -= 1,
            _ => (),
        }
        count += 1;
        if level == 0 {
            break;
        }
    }

    count == tokens.len()
}

pub fn build_ast(mut tokens: Vec<Token>) -> Node {
    //let mut tokens_copy = tokens.clone();
    //let mut not_done = true;

    // Lowest prio : equal
    let split = tokens.iter().position(|x| x == &Token::Equal());
    if let Some(position) = split {
        return insert_operator_node(tokens, position);
    }

    while tokens.starts_with(&[Token::OpenParenthesis()]) && tokens.ends_with(&[Token::CloseParenthesis()]) && check_parentheses(&tokens) == true {
        tokens = tokens.split_off(1);
        tokens.pop();
    }

    let parentheses_map = get_parentheses_map(&tokens);
    //let iter = zip(tokens, parentheses_map);
    // Lowest prio : rightest minus or plus, also manages minus unary operator
    let split = tokens.iter().enumerate().rposition(|(i, x)| parentheses_map[i] == 0 && (x == &Token::Addition() || x == &Token::Substraction()));
    if let Some(position) = split {
        let tokens_right = tokens.split_off(position + 1);
        let token = tokens.pop().unwrap();
        if tokens.len() == 0 {
            match &tokens_right[0] {
                Token::Number(n) => return Node::Leaf(vec![parse_number(n, -1)]),
                Token::ImaginaryUnit() => return Node::Leaf(vec![minus_i()]),
                Token::Variable(_) => return Node::Leaf(minus_x()),
                _ => (),
            }
        }
        return Node::Internal {
            operator: token,
            lhs: Box::new(build_ast(tokens)),
            rhs: Box::new(build_ast(tokens_right)),
        }
    }

    // Lowest prio : righest times or divide or modulo
    let split = tokens.iter().enumerate().rposition(|(i, x)| parentheses_map[i] == 0 && (x == &Token::Multiplication() || x == &Token::Division() || x == &Token::Modulo()));
    if let Some(position) = split {
        return insert_operator_node(tokens, position);
    }
    // Lowest prio : exponentiation
    let split = tokens.iter().enumerate().rposition(|(i, x)| parentheses_map[i] == 0 && x == &Token::Exponentiation());
    if let Some(position) = split {
        return insert_operator_node(tokens, position);
    }

    match &tokens[0] {
        Token::Number(n) => return Node::Leaf(vec![parse_number(n, 1)]),
        Token::ImaginaryUnit() => return Node::Leaf(vec![i()]),
        Token::Variable(_) => return Node::Leaf(x()),
        _ => (),
    }
    unreachable!();
}

#[cfg(test)]
mod tests {
    use super::*;

    fn one_token() -> Token {
        Token::Number(String::from("1"))
    }

    fn two_token() -> Token {
        Token::Number(String::from("2"))
    }

    fn x_token() -> Token {
        Token::Variable(String::from("x"))
    }

    fn i_token() -> Token {
        Token::ImaginaryUnit()
    }

    fn open() -> Token {
        Token::OpenParenthesis()
    }

    fn close() -> Token {
        Token::CloseParenthesis()
    }

    fn plus() -> Token {
        Token::Addition()
    }

    fn minus() -> Token {
        Token::Substraction()
    }

    fn times() -> Token {
        Token::Multiplication()
    }

    fn divided_by() -> Token {
        Token::Division()
    }

    fn modulo() -> Token {
        Token::Modulo()
    }

    fn power() -> Token {
        Token::Exponentiation()
    }

    fn equals() -> Token {
        Token::Equal()
    }

    #[test]
    fn one_equals_one() {
        let tokens = vec![one_token(), equals(), one_token()];
        let results = Node::Internal {
            operator: Token::Equal(),
            lhs: Box::new(Node::Leaf(vec![one()])),
            rhs: Box::new(Node::Leaf(vec![one()]))
        };

        assert_eq!(build_ast(tokens), results);

    }

    #[test]
    fn one_plus_one() {
        let tokens = vec![one_token(), plus(), one_token()];
        let results = Node::Internal {
            operator: plus(),
            lhs: Box::new(Node::Leaf(vec![one()])),
            rhs: Box::new(Node::Leaf(vec![one()]))
        };

        assert_eq!(build_ast(tokens), results);

        let tokens = vec![one_token(), minus(), one_token()];
        let results = Node::Internal {
            operator: minus(),
            lhs: Box::new(Node::Leaf(vec![one()])),
            rhs: Box::new(Node::Leaf(vec![one()]))
        };

        assert_eq!(build_ast(tokens), results);
    }

    #[test]
    fn one_times_one() {
        let tokens = vec![one_token(), times(), one_token()];
        let results = Node::Internal {
            operator: times(),
            lhs: Box::new(Node::Leaf(vec![one()])),
            rhs: Box::new(Node::Leaf(vec![one()]))
        };
        assert_eq!(build_ast(tokens), results);

        let tokens = vec![one_token(), divided_by(), one_token()];
        let results = Node::Internal {
            operator: divided_by(),
            lhs: Box::new(Node::Leaf(vec![one()])),
            rhs: Box::new(Node::Leaf(vec![one()]))
        };
        assert_eq!(build_ast(tokens), results);

        let tokens = vec![one_token(), modulo(), one_token()];
        let results = Node::Internal {
            operator: modulo(),
            lhs: Box::new(Node::Leaf(vec![one()])),
            rhs: Box::new(Node::Leaf(vec![one()]))
        };
        assert_eq!(build_ast(tokens), results);
    }

    #[test]
    fn one_power_one() {
        let tokens = vec![one_token(), power(), one_token()];
        let results = Node::Internal {
            operator: power(),
            lhs: Box::new(Node::Leaf(vec![one()])),
            rhs: Box::new(Node::Leaf(vec![one()]))
        };
        assert_eq!(build_ast(tokens), results);
    }

    #[test]
    fn left_priority() {
        let tokens = vec![one_token(), plus(), one_token(), minus(), one_token()];
        let results = Node::Internal {
            operator: minus(),
            lhs: Box::new(Node::Internal {
                operator: plus(),
                lhs: Box::new(Node::Leaf(vec![one()])),
                rhs: Box::new(Node::Leaf(vec![one()]))
            }),
            rhs: Box::new(Node::Leaf(vec![one()]))
        };
        assert_eq!(build_ast(tokens), results);
    }

    #[test]
    fn operation_priority() {
        let tokens = vec![one_token(), plus(), one_token(), times(), one_token()];
        let results = Node::Internal {
            operator: plus(),
            lhs: Box::new(Node::Leaf(vec![one()])),
            rhs: Box::new(Node::Internal {
                operator: times(),
                lhs: Box::new(Node::Leaf(vec![one()])),
                rhs: Box::new(Node::Leaf(vec![one()]))
            })
        };
        assert_eq!(build_ast(tokens), results);
    }

    #[test]
    fn i_plus_i() {
        let tokens = vec![i_token(), plus(), i_token()];
        let results = Node::Internal {
            operator: plus(),
            lhs: Box::new(Node::Leaf(vec![i()])),
            rhs: Box::new(Node::Leaf(vec![i()]))
        };

        assert_eq!(build_ast(tokens), results);
    }

    #[test]
    fn x_plus_x() {
        let tokens = vec![x_token(), plus(), x_token()];
        let results = Node::Internal {
            operator: plus(),
            lhs: Box::new(Node::Leaf(x())),
            rhs: Box::new(Node::Leaf(x()))
        };

        assert_eq!(build_ast(tokens), results);
    }

    #[test]
    fn parentheses() {
        let tokens = vec![one_token(), times(), open(), one_token(), plus(), one_token(), close()];
        let results = Node::Internal {
            operator: times(),
            lhs: Box::new(Node::Leaf(vec![one()])),
            rhs: Box::new(Node::Internal {
                operator: plus(),
                lhs: Box::new(Node::Leaf(vec![one()])),
                rhs: Box::new(Node::Leaf(vec![one()]))
            })
        };
        assert_eq!(build_ast(tokens), results);
    }

    #[test]
    fn two_plus_two() {
        let mut two = one();
        two.real.numerator = 2;

        let tokens = vec![two_token(), plus(), two_token()];
        let results = Node::Internal {
            operator: plus(),
            lhs: Box::new(Node::Leaf(vec![two.clone()])),
            rhs: Box::new(Node::Leaf(vec![two.clone()]))
        };

        assert_eq!(build_ast(tokens), results);

        let tokens = vec![two_token(), minus(), one_token()];
        let results = Node::Internal {
            operator: minus(),
            lhs: Box::new(Node::Leaf(vec![two])),
            rhs: Box::new(Node::Leaf(vec![one()]))
        };

        assert_eq!(build_ast(tokens), results);
    }

    #[test]
    fn negative_one() {
        let tokens = vec![minus(), one_token()];
        let results = Node::Leaf(vec![minus_one()]);
        assert_eq!(build_ast(tokens), results);

        let tokens = vec![minus(), x_token()];
        let results = Node::Leaf(minus_x());
        assert_eq!(build_ast(tokens), results);

        let tokens = vec![minus(), i_token()];
        let results = Node::Leaf(vec![minus_i()]);
        assert_eq!(build_ast(tokens), results);
    }
}