test_engine.py

# -*- coding: utf-8 -*-
"""
独立测试脚本：验证 MathFunction 升级版表达式引擎与数学运算
"""
import math
import re

# 1. 表达式节点定义
class _Num:
    __slots__ = ("v",)
    def __init__(self, v): self.v = float(v)

class _Var:
    __slots__ = ("name",)
    def __init__(self, name): self.name = name

class _Un:
    __slots__ = ("op", "child")
    def __init__(self, op, child): self.op = op; self.child = child

class _Bin:
    __slots__ = ("op", "l", "r")
    def __init__(self, op, l, r): self.op = op; self.l = l; self.r = r

class _Call:
    __slots__ = ("name", "args")
    def __init__(self, name, args): self.name = name; self.args = args

# 2. 安全数学运算
def safe_div(a, b):
    try:
        if abs(b) < 1e-14:
            return float('nan')
        return a / b
    except Exception:
        return float('nan')

def safe_sqrt(x):
    if x < 0:
        return float('nan')
    return math.sqrt(x)

def safe_log(x):
    if x <= 0:
        return float('nan')
    return math.log(x)

def safe_log10(x):
    if x <= 0:
        return float('nan')
    return math.log10(x)

def safe_log2(x):
    if x <= 0:
        return float('nan')
    return math.log2(x)

def safe_pow(a, b):
    try:
        if a < 0 and int(b) != b:
            return float('nan')
        res = a ** b
        if abs(res) > 1e15:
            return math.copysign(1e15, res)
        return res
    except Exception:
        return float('nan')

def safe_tan(x):
    try:
        c = math.cos(x)
        if abs(c) < 1e-12:
            return float('nan')
        return math.sin(x) / c
    except Exception:
        return float('nan')

def safe_sinc(x):
    if abs(x) < 1e-7:
        return 1.0 - (x * x) / 6.0
    return math.sin(x) / x

def safe_clamp(x, min_v, max_v):
    return max(min_v, min(max_v, x))

def safe_step(edge, x):
    return 1.0 if x >= edge else 0.0

def safe_smoothstep(edge0, edge1, x):
    if abs(edge1 - edge0) < 1e-12:
        return 0.0
    t = safe_clamp((x - edge0) / (edge1 - edge0), 0.0, 1.0)
    return t * t * (3.0 - 2.0 * t)

def safe_lerp(a, b, t):
    return a + (b - a) * t

def safe_if(cond, a, b):
    return a if bool(cond) else b

def safe_cbrt(x):
    return math.copysign(abs(x) ** (1.0 / 3.0), x)

_FN = {}
def _reg_fn(name, fn, deriv=None):
    _FN[name] = (fn, deriv)

_reg_fn("sin", math.sin, "cos")
_reg_fn("cos", math.cos, "sin")
_reg_fn("tan", safe_tan, "tan")
_reg_fn("asin", lambda x: math.asin(safe_clamp(x, -1.0, 1.0)), "asin")
_reg_fn("acos", lambda x: math.acos(safe_clamp(x, -1.0, 1.0)), "acos")
_reg_fn("atan", math.atan, "atan")
_reg_fn("sinh", math.sinh, "cosh")
_reg_fn("cosh", math.cosh, "sinh")
_reg_fn("tanh", math.tanh, "tanh")
_reg_fn("sqrt", safe_sqrt, "sqrt")
_reg_fn("cbrt", safe_cbrt, None)
_reg_fn("exp", lambda x: math.exp(min(x, 700.0)), "exp")
_reg_fn("ln", safe_log, "ln")
_reg_fn("log", safe_log, "ln")
_reg_fn("log10", safe_log10, "log10")
_reg_fn("log2", safe_log2, "log2")
_reg_fn("abs", abs, "abs")
_reg_fn("sign", lambda x: (1.0 if x > 0 else (-1.0 if x < 0 else 0.0)), None)
_reg_fn("sgn", lambda x: (1.0 if x > 0 else (-1.0 if x < 0 else 0.0)), None)
_reg_fn("floor", math.floor, None)
_reg_fn("ceil", math.ceil, None)
_reg_fn("round", round, None)
_reg_fn("trunc", math.trunc, None)
_reg_fn("sec", lambda x: safe_div(1.0, math.cos(x)), "sec")
_reg_fn("csc", lambda x: safe_div(1.0, math.sin(x)), "csc")
_reg_fn("cot", lambda x: safe_div(1.0, safe_tan(x)), "cot")
_reg_fn("gamma", math.gamma, None)
_reg_fn("erf", math.erf, "erf")
_reg_fn("degrees", math.degrees, None)
_reg_fn("radians", math.radians, None)
_reg_fn("pow", safe_pow, None)
_reg_fn("atan2", math.atan2, None)
_reg_fn("hypot", math.hypot, None)
_reg_fn("min", min, None)
_reg_fn("max", max, None)
_reg_fn("sinc", safe_sinc, "sinc")
_reg_fn("clamp", safe_clamp, None)
_reg_fn("step", safe_step, None)
_reg_fn("smoothstep", safe_smoothstep, None)
_reg_fn("lerp", safe_lerp, None)
_reg_fn("mod", lambda a, b: a % b if abs(b) > 1e-12 else float('nan'), None)
_reg_fn("if", safe_if, None)

_CONSTANTS = {
    "pi": math.pi,
    "e": math.e,
    "tau": math.tau,
    "phi": (1.0 + math.sqrt(5.0)) / 2.0,
    "inf": float('inf'),
}

# 3. 分词与隐式乘法处理
_TOK_RE = re.compile(
    r"\s*(?:(\d+\.\d*|\.\d+|\d+)(?:[eE][-+]?\d+)?|([a-zA-Z_][a-zA-Z_0-9]*)|(==|!=|<=|>=|\*\*|[+\-*/^(),<>=?:;]))"
)

def _raw_tokenize(s):
    toks = []
    i = 0
    n = len(s)
    while i < n:
        m = _TOK_RE.match(s, i)
        if not m:
            raise ValueError(f"无法识别的字符: '{s[i]}' (位置 {i})")
        if m.group(1) is not None:
            toks.append(("num", float(m.group(0))))
        elif m.group(2):
            toks.append(("id", m.group(2)))
        elif m.group(3):
            toks.append(("op", m.group(3)))
        i = m.end()
    return toks

def _insert_implicit_mul(raw_toks):
    """
    在需要隐式乘法的位置插入 ('op', '*')：
    例如：
    - 2 x       -> 2 * x
    - 2 sin(x)  -> 2 * sin(x)
    - 2 (x+1)   -> 2 * (x+1)
    - (x+1)(x-1)-> (x+1) * (x-1)
    - (x+1) x   -> (x+1) * x
    - x (x+1)   (若 x 不是已注册函数名) -> x * (x+1)
    """
    out = []
    n = len(raw_toks)
    for i in range(n):
        curr = raw_toks[i]
        out.append(curr)
        if i + 1 < n:
            nxt = raw_toks[i + 1]
            c_type, c_val = curr
            n_type, n_val = nxt

            # 前一个 token 可以是数字、括号右、或变量（非函数名紧邻括号）
            can_end = False
            if c_type == "num":
                can_end = True
            elif c_type == "op" and c_val == ")":
                can_end = True
            elif c_type == "id":
                # 只有当它是变量或常量，且后面不是紧跟 '(' 时（或者是紧跟 '(' 但它不是已知函数）
                if not (n_type == "op" and n_val == "(" and c_val in _FN):
                    can_end = True

            # 后一个 token 可以是数字、变量/函数、或括号左
            can_start = False
            if n_type == "num":
                can_start = True
            elif n_type == "id":
                can_start = True
            elif n_type == "op" and n_val == "(":
                # 如果前一个不是函数名
                if not (c_type == "id" and c_val in _FN):
                    can_start = True

            if can_end and can_start:
                out.append(("op", "*"))

    out.append(("end", None))
    return out

def _tokenize(s):
    return _insert_implicit_mul(_raw_tokenize(s))

# 4. 语法解析器 (递归下降)
class _Parser:
    def __init__(self, toks):
        self.toks = toks
        self.i = 0

    def peek(self):
        return self.toks[self.i]

    def next(self):
        t = self.toks[self.i]
        self.i += 1
        return t

    def parse(self):
        e = self.expr()
        if self.peek()[0] != "end":
            raise ValueError(f"表达式结尾有多余内容: {self.peek()[1]!r}")
        return e

    def expr(self):
        # 比较运算与三元条件支持
        return self.ternary()

    def ternary(self):
        node = self.comparison()
        if self.peek()[0] == "op" and self.peek()[1] == "?":
            self.next()
            true_expr = self.expr()
            if self.peek()[0] != "op" or self.peek()[1] != ":":
                raise ValueError("三元表达式缺少 ':'")
            self.next()
            false_expr = self.expr()
            return _Call("if", [node, true_expr, false_expr])
        return node

    def comparison(self):
        node = self.additive()
        while self.peek()[0] == "op" and self.peek()[1] in ("==", "!=", "<", ">", "<=", ">="):
            op = self.next()[1]
            r = self.additive()
            node = _Bin(op, node, r)
        return node

    def additive(self):
        node = self.term()
        while self.peek()[0] == "op" and self.peek()[1] in ("+", "-"):
            op = self.next()[1]
            r = self.term()
            node = _Bin(op, node, r)
        return node

    def term(self):
        node = self.factor()
        while self.peek()[0] == "op" and self.peek()[1] in ("*", "/"):
            op = self.next()[1]
            r = self.factor()
            node = _Bin(op, node, r)
        return node

    def factor(self):
        t = self.peek()
        if t[0] == "op" and t[1] in ("+", "-"):
            self.next()
            return _Un(t[1], self.factor())
        return self.power()

    def power(self):
        node = self.primary()
        if self.peek()[0] == "op" and self.peek()[1] in ("^", "**"):
            self.next()
            r = self.factor()  # 右结合
            node = _Bin("^", node, r)
        return node

    def primary(self):
        t = self.next()
        if t[0] == "num":
            return _Num(t[1])
        if t[0] == "id":
            name = t[1]
            if name in _CONSTANTS:
                return _Num(_CONSTANTS[name])
            if self.peek()[0] == "op" and self.peek()[1] == "(":
                self.next()  # '('
                args = []
                if not (self.peek()[0] == "op" and self.peek()[1] == ")"):
                    args.append(self.expr())
                    while self.peek()[0] == "op" and self.peek()[1] == ",":
                        self.next()
                        args.append(self.expr())
                if self.peek()[0] != "op" or self.peek()[1] != ")":
                    raise ValueError(f"函数 '{name}' 调用缺少右括号 ')'")
                self.next()  # ')'
                if name not in _FN:
                    raise ValueError(f"未知函数: '{name}'")
                return _Call(name, args)
            return _Var(name)
        if t[0] == "op" and t[1] == "(":
            node = self.expr()
            if self.peek()[0] != "op" or self.peek()[1] != ")":
                raise ValueError("缺少右括号 ')'")
            self.next()
            return node
        raise ValueError(f"意外的符号: {t[1]!r}")

def parse(expr):
    return _Parser(_tokenize(expr)).parse()

def _compile(node):
    if isinstance(node, _Num):
        v = node.v
        return lambda env: v
    if isinstance(node, _Var):
        n = node.name
        return lambda env: env.get(n, 0.0)
    if isinstance(node, _Un):
        c = _compile(node.child)
        if node.op == "-":
            return lambda env: -c(env)
        return c
    if isinstance(node, _Bin):
        a = _compile(node.l)
        b = _compile(node.r)
        op = node.op
        if op == "+":
            return lambda env: a(env) + b(env)
        if op == "-":
            return lambda env: a(env) - b(env)
        if op == "*":
            return lambda env: a(env) * b(env)
        if op == "/":
            return lambda env: safe_div(a(env), b(env))
        if op in ("^", "**"):
            return lambda env: safe_pow(a(env), b(env))
        if op == "==":
            return lambda env: 1.0 if abs(a(env) - b(env)) < 1e-9 else 0.0
        if op == "!=":
            return lambda env: 1.0 if abs(a(env) - b(env)) >= 1e-9 else 0.0
        if op == "<":
            return lambda env: 1.0 if a(env) < b(env) else 0.0
        if op == ">":
            return lambda env: 1.0 if a(env) > b(env) else 0.0
        if op == "<=":
            return lambda env: 1.0 if a(env) <= b(env) else 0.0
        if op == ">=":
            return lambda env: 1.0 if a(env) >= b(env) else 0.0
        raise ValueError(f"未知运算符: {op}")
    if isinstance(node, _Call):
        fn = _FN[node.name][0]
        args = [_compile(a) for a in node.args]
        return lambda env: fn(*[a(env) for a in args])
    raise TypeError(node)

def compile_expression(expr):
    ast = parse(expr.strip())
    return _compile(ast), ast

# 5. 测试套件
def run_tests():
    test_cases = [
        ("2x", {"x": 3}, 6.0),
        ("2(x+1)", {"x": 3}, 8.0),
        ("2sin(pi/2)", {}, 2.0),
        ("x**2 + 1", {"x": 4}, 17.0),
        ("x^3", {"x": 2}, 8.0),
        ("(x+1)(x-1)", {"x": 5}, 24.0),
        ("1/x", {"x": 0.0}, float('nan')),  # safe_div test
        ("sqrt(-4)", {}, float('nan')),      # safe_sqrt test
        ("sinc(0)", {}, 1.0),
        ("if(x > 0, 10, -10)", {"x": 5}, 10.0),
        ("if(x > 0, 10, -10)", {"x": -2}, -10.0),
        ("x > 0 ? 10 : -10", {"x": 5}, 10.0),
        ("clamp(x, 0, 1)", {"x": 1.5}, 1.0),
        ("smoothstep(0, 1, 0.5)", {}, 0.5),
        ("sin(x)*cos(y)", {"x": 0, "y": 0}, 0.0),
    ]

    print("Running engine unit tests...")
    all_ok = True
    for expr, env, expected in test_cases:
        try:
            fn, _ = compile_expression(expr)
            val = fn(env)
            if math.isnan(expected):
                ok = math.isnan(val)
            else:
                ok = abs(val - expected) < 1e-6
            status = "PASS" if ok else f"FAIL (got {val}, expected {expected})"
            print(f"  [{status}] '{expr}' with {env} -> {val}")
            if not ok:
                all_ok = False
        except Exception as e:
            print(f"  [ERROR] '{expr}' failed with exception: {e}")
            all_ok = False

    if all_ok:
        print("\nAll unit tests passed successfully!")
    else:
        print("\nSome tests failed!")
    return all_ok

if __name__ == "__main__":
    run_tests()