ImHex/plugins/builtin/source/math_evaluator.cpp

#include "math_evaluator.hpp"

#include <string>
#include <queue>
#include <stack>
#include <stdexcept>
#include <cmath>
#include <cstdint>
#include <optional>
#include <numbers>

namespace hex {

    i16 comparePrecedence(const Operator &a, const Operator &b) {
        return (static_cast<i8>(a) & 0x0F0) - (static_cast<i8>(b) & 0x0F0);
    }

    bool isLeftAssociative(const Operator op) {
        return (static_cast<u32>(op) & 0xF00) == 0;
    }

    std::pair<Operator, size_t> toOperator(std::string input) {
        if (input.starts_with("##")) return { Operator::Combine, 2 };
        if (input.starts_with("==")) return { Operator::Equals, 2 };
        if (input.starts_with("!=")) return { Operator::NotEquals, 2 };
        if (input.starts_with(">=")) return { Operator::GreaterThanOrEquals, 2 };
        if (input.starts_with("<=")) return { Operator::LessThanOrEquals, 2 };
        if (input.starts_with(">>")) return { Operator::ShiftRight, 2 };
        if (input.starts_with("<<")) return { Operator::ShiftLeft, 2 };
        if (input.starts_with("||")) return { Operator::Or, 2 };
        if (input.starts_with("^^")) return { Operator::Xor, 2 };
        if (input.starts_with("&&")) return { Operator::And, 2 };
        if (input.starts_with("**")) return { Operator::Exponentiation, 2 };
        if (input.starts_with(">")) return { Operator::GreaterThan, 1 };
        if (input.starts_with("<")) return { Operator::LessThan, 1 };
        if (input.starts_with("!")) return { Operator::Not, 1 };
        if (input.starts_with("|")) return { Operator::BitwiseOr, 1 };
        if (input.starts_with("^")) return { Operator::BitwiseXor, 1 };
        if (input.starts_with("&")) return { Operator::BitwiseAnd, 1 };
        if (input.starts_with("~")) return { Operator::BitwiseNot, 1 };
        if (input.starts_with("+")) return { Operator::Addition, 1 };
        if (input.starts_with("-")) return { Operator::Subtraction, 1 };
        if (input.starts_with("*")) return { Operator::Multiplication, 1 };
        if (input.starts_with("/")) return { Operator::Division, 1 };
        if (input.starts_with("%")) return { Operator::Modulus, 1 };
        if (input.starts_with("=")) return { Operator::Assign, 1 };

        return { Operator::Invalid, 0 };
    }

    std::queue<Token> MathEvaluator::parseInput(const char *input) {
        std::queue<Token> inputQueue;

        char *prevPos = const_cast<char *>(input);
        for (char *pos = prevPos; *pos != 0x00;) {
            if (std::isdigit(*pos) || *pos == '.') {
                auto number = std::strtold(pos, &pos);

                if (*pos == 'x') {
                    pos--;
                    number = std::strtoull(pos, &pos, 0);
                }

                inputQueue.push(Token { .type = TokenType::Number, .number = number });
            } else if (*pos == '(') {
                inputQueue.push(Token { .type = TokenType::Bracket, .bracketType = BracketType::Left });
                pos++;
            } else if (*pos == ')') {
                inputQueue.push(Token { .type = TokenType::Bracket, .bracketType = BracketType::Right });
                pos++;
            } else if (std::isspace(*pos)) {
                pos++;
            } else {
                auto [op, width] = toOperator(pos);

                if (op != Operator::Invalid) {
                    inputQueue.push(Token { .type = TokenType::Operator, .op = op });
                    pos += width;
                } else {
                    Token token;

                    while (std::isalpha(*pos) || *pos == '_') {
                        token.name += *pos;
                        pos++;
                    }

                    if (*pos == '(') {
                        pos++;

                        u32 depth = 1;
                        std::vector<std::string> expressions;
                        expressions.emplace_back();

                        while (*pos != 0x00) {
                            if (*pos == '(') depth++;
                            else if (*pos == ')') depth--;

                            if (depth == 0)
                                break;

                            if (depth == 1 && *pos == ',') {
                                expressions.emplace_back();
                                pos++;
                            }

                            expressions.back() += *pos;

                            pos++;
                        }

                        pos++;

                        for (const auto &expression : expressions) {
                            if (expression == "" && expressions.size() > 1)
                                throw std::invalid_argument("Invalid function call syntax!");
                            else if (expression == "")
                                break;

                            auto inputQueue = parseInput(expression.c_str());
                            auto postfixTokens = toPostfix(inputQueue);
                            auto result = evaluate(postfixTokens);

                            if (!result.has_value())
                                throw std::invalid_argument("Invalid argument for function!");

                            token.arguments.push_back(result.value());
                        }

                        token.type = TokenType::Function;
                        inputQueue.push(token);

                    } else {
                        token.type = TokenType::Variable;
                        inputQueue.push(token);
                    }
                }
            }

            if (prevPos == pos)
                throw std::invalid_argument("Invalid syntax!");

            prevPos = pos;
        }

        return inputQueue;
    }

    std::queue<Token> MathEvaluator::toPostfix(std::queue<Token> inputQueue) {
        std::queue<Token> outputQueue;
        std::stack<Token> operatorStack;

        while (!inputQueue.empty()) {
            Token currToken = inputQueue.front();
            inputQueue.pop();

            if (currToken.type == TokenType::Number || currToken.type == TokenType::Variable || currToken.type == TokenType::Function)
                outputQueue.push(currToken);
            else if (currToken.type == TokenType::Operator) {
                while ((!operatorStack.empty()) && (operatorStack.top().type == TokenType::Operator && currToken.type == TokenType::Operator && (comparePrecedence(operatorStack.top().op, currToken.op) > 0) || (comparePrecedence(operatorStack.top().op, currToken.op) == 0 && isLeftAssociative(currToken.op))) && operatorStack.top().type != TokenType::Bracket) {
                    outputQueue.push(operatorStack.top());
                    operatorStack.pop();
                }
                operatorStack.push(currToken);
            } else if (currToken.type == TokenType::Bracket) {
                if (currToken.bracketType == BracketType::Left)
                    operatorStack.push(currToken);
                else {
                    if (operatorStack.empty())
                        throw std::invalid_argument("Mismatching parenthesis!");

                    while (operatorStack.top().type != TokenType::Bracket || (operatorStack.top().type == TokenType::Bracket && operatorStack.top().bracketType != BracketType::Left)) {
                        if (operatorStack.empty())
                            throw std::invalid_argument("Mismatching parenthesis!");

                        outputQueue.push(operatorStack.top());
                        operatorStack.pop();
                    }

                    operatorStack.pop();
                }
            }
        }

        while (!operatorStack.empty()) {
            auto top = operatorStack.top();

            if (top.type == TokenType::Bracket)
                throw std::invalid_argument("Mismatching parenthesis!");

            outputQueue.push(top);
            operatorStack.pop();
        }

        return outputQueue;
    }

    std::optional<long double> MathEvaluator::evaluate(std::queue<Token> postfixTokens) {
        std::stack<long double> evaluationStack;

        while (!postfixTokens.empty()) {
            auto front = postfixTokens.front();
            postfixTokens.pop();

            if (front.type == TokenType::Number)
                evaluationStack.push(front.number);
            else if (front.type == TokenType::Operator) {
                long double rightOperand, leftOperand;
                if (evaluationStack.size() < 2) {
                    if ((front.op == Operator::Addition || front.op == Operator::Subtraction || front.op == Operator::Not || front.op == Operator::BitwiseNot) && evaluationStack.size() == 1) {
                        rightOperand = evaluationStack.top();
                        evaluationStack.pop();
                        leftOperand = 0;
                    } else throw std::invalid_argument("Not enough operands for operator!");
                } else {
                    rightOperand = evaluationStack.top();
                    evaluationStack.pop();
                    leftOperand = evaluationStack.top();
                    evaluationStack.pop();
                }

                long double result = std::numeric_limits<long double>::quiet_NaN();
                switch (front.op) {
                default:
                case Operator::Invalid:
                    throw std::invalid_argument("Invalid operator!");
                case Operator::And:
                    result = static_cast<i64>(leftOperand) && static_cast<i64>(rightOperand);
                    break;
                case Operator::Or:
                    result = static_cast<i64>(leftOperand) && static_cast<i64>(rightOperand);
                    break;
                case Operator::Xor:
                    result = (static_cast<i64>(leftOperand) ^ static_cast<i64>(rightOperand)) > 0;
                    break;
                case Operator::GreaterThan:
                    result = leftOperand > rightOperand;
                    break;
                case Operator::LessThan:
                    result = leftOperand < rightOperand;
                    break;
                case Operator::GreaterThanOrEquals:
                    result = leftOperand >= rightOperand;
                    break;
                case Operator::LessThanOrEquals:
                    result = leftOperand <= rightOperand;
                    break;
                case Operator::Equals:
                    result = leftOperand == rightOperand;
                    break;
                case Operator::NotEquals:
                    result = leftOperand != rightOperand;
                    break;
                case Operator::Not:
                    result = !static_cast<i64>(rightOperand);
                    break;
                case Operator::BitwiseOr:
                    result = static_cast<i64>(leftOperand) | static_cast<i64>(rightOperand);
                    break;
                case Operator::BitwiseXor:
                    result = static_cast<i64>(leftOperand) ^ static_cast<i64>(rightOperand);
                    break;
                case Operator::BitwiseAnd:
                    result = static_cast<i64>(leftOperand) & static_cast<i64>(rightOperand);
                    break;
                case Operator::BitwiseNot:
                    result = ~static_cast<i64>(rightOperand);
                    break;
                case Operator::ShiftLeft:
                    result = static_cast<i64>(leftOperand) << static_cast<i64>(rightOperand);
                    break;
                case Operator::ShiftRight:
                    result = static_cast<i64>(leftOperand) >> static_cast<i64>(rightOperand);
                    break;
                case Operator::Addition:
                    result = leftOperand + rightOperand;
                    break;
                case Operator::Subtraction:
                    result = leftOperand - rightOperand;
                    break;
                case Operator::Multiplication:
                    result = leftOperand * rightOperand;
                    break;
                case Operator::Division:
                    result = leftOperand / rightOperand;
                    break;
                case Operator::Modulus:
                    result = std::fmod(leftOperand, rightOperand);
                    break;
                case Operator::Exponentiation:
                    result = std::pow(leftOperand, rightOperand);
                    break;
                case Operator::Combine:
                    result = (static_cast<u64>(leftOperand) << (64 - __builtin_clzll(static_cast<u64>(rightOperand)))) | static_cast<u64>(rightOperand);
                    break;
                }

                evaluationStack.push(result);
            } else if (front.type == TokenType::Variable) {
                if (this->m_variables.contains(front.name))
                    evaluationStack.push(this->m_variables.at(front.name));
                else
                    throw std::invalid_argument("Unknown variable!");
            } else if (front.type == TokenType::Function) {
                if (!this->m_functions[front.name])
                    throw std::invalid_argument("Unknown function called!");

                auto result = this->m_functions[front.name](front.arguments);

                if (result.has_value())
                    evaluationStack.push(result.value());
            } else
                throw std::invalid_argument("Parenthesis in postfix expression!");
        }

        if (evaluationStack.empty())
            return {};
        else if (evaluationStack.size() > 1)
            throw std::invalid_argument("Undigested input left!");
        else
            return evaluationStack.top();
    }


    std::optional<long double> MathEvaluator::evaluate(std::string input) {
        auto inputQueue = parseInput(input.c_str());

        std::string resultVariable = "ans";

        {
            std::queue<Token> queueCopy = inputQueue;
            if (queueCopy.front().type == TokenType::Variable) {
                resultVariable = queueCopy.front().name;
                queueCopy.pop();
                if (queueCopy.front().type != TokenType::Operator || queueCopy.front().op != Operator::Assign)
                    resultVariable = "ans";
                else {
                    inputQueue.pop();
                    inputQueue.pop();
                }
            }
        }

        auto postfixTokens = toPostfix(inputQueue);

        auto result = evaluate(postfixTokens);

        if (result.has_value()) {
            this->setVariable(resultVariable, result.value());
        }

        return result;
    }

    void MathEvaluator::setVariable(std::string name, long double value) {
        this->m_variables[name] = value;
    }

    void MathEvaluator::setFunction(std::string name, std::function<std::optional<long double>(std::vector<long double>)> function, size_t minNumArgs, size_t maxNumArgs) {
        this->m_functions[name] = [minNumArgs, maxNumArgs, function](auto args) {
            if (args.size() < minNumArgs || args.size() > maxNumArgs)
                throw std::invalid_argument("Invalid number of function arguments!");

            return function(args);
        };
    }


    void MathEvaluator::registerStandardVariables() {
        this->setVariable("ans", 0);
    }

    void MathEvaluator::registerStandardFunctions() {
        this->setFunction(
            "sin", [](auto args) { return std::sin(args[0]); }, 1, 1);
        this->setFunction(
            "cos", [](auto args) { return std::cos(args[0]); }, 1, 1);
        this->setFunction(
            "tan", [](auto args) { return std::tan(args[0]); }, 1, 1);
        this->setFunction(
            "sqrt", [](auto args) { return std::sqrt(args[0]); }, 1, 1);
        this->setFunction(
            "ceil", [](auto args) { return std::ceil(args[0]); }, 1, 1);
        this->setFunction(
            "floor", [](auto args) { return std::floor(args[0]); }, 1, 1);
        this->setFunction(
            "sign", [](auto args) { return (args[0] > 0) ? 1 : (args[0] == 0) ? 0
                                                                              : -1; }, 1, 1);
        this->setFunction(
            "abs", [](auto args) { return std::abs(args[0]); }, 1, 1);
        this->setFunction(
            "ln", [](auto args) { return std::log(args[0]); }, 1, 1);
        this->setFunction(
            "lb", [](auto args) { return std::log2(args[0]); }, 1, 1);
        this->setFunction(
            "log", [](auto args) { return args.size() == 1 ? std::log10(args[0]) : std::log(args[1]) / std::log(args[0]); }, 1, 2);
    }

}