Evaluating a mathematical [removed]function) for a large number of input values fast

前端 未结 5 1635
你的背包
你的背包 2020-12-09 12:11

The following questions

  • Evaluating a mathematical expression in a string
  • Equation parsing in Python
  • Safe way to parse user-supplied
5条回答
  •  慢半拍i
    慢半拍i (楼主)
    2020-12-09 12:40

    CPython (and pypy) use a very simple stack language for executing functions, and it is fairly easy to write the bytecode yourself, using the ast module.

    import sys
    PY3 = sys.version_info.major > 2
    import ast
    from ast import parse
    import types
    from dis import opmap
    
    ops = {
        ast.Mult: opmap['BINARY_MULTIPLY'],
        ast.Add: opmap['BINARY_ADD'],
        ast.Sub: opmap['BINARY_SUBTRACT'],
        ast.Div: opmap['BINARY_TRUE_DIVIDE'],
        ast.Pow: opmap['BINARY_POWER'],
    }
    LOAD_CONST = opmap['LOAD_CONST']
    RETURN_VALUE = opmap['RETURN_VALUE']
    LOAD_FAST = opmap['LOAD_FAST']
    def process(consts, bytecode, p, stackSize=0):
        if isinstance(p, ast.Expr):
            return process(consts, bytecode, p.value, stackSize)
        if isinstance(p, ast.BinOp):
            szl = process(consts, bytecode, p.left, stackSize)
            szr = process(consts, bytecode, p.right, stackSize)
            if type(p.op) in ops:
                bytecode.append(ops[type(p.op)])
            else:
                print(p.op)
                raise Exception("unspported opcode")
            return max(szl, szr) + stackSize + 1
        if isinstance(p, ast.Num):
            if p.n not in consts:
                consts.append(p.n)
            idx = consts.index(p.n)
            bytecode.append(LOAD_CONST)
            bytecode.append(idx % 256)
            bytecode.append(idx // 256)
            return stackSize + 1
        if isinstance(p, ast.Name):
            bytecode.append(LOAD_FAST)
            bytecode.append(0)
            bytecode.append(0)
            return stackSize + 1
        raise Exception("unsupported token")
    
    def makefunction(inp):
        def f(x):
            pass
    
        if PY3:
            oldcode = f.__code__
            kwonly = oldcode.co_kwonlyargcount
        else:
            oldcode = f.func_code
        stack_size = 0
        consts = [None]
        bytecode = []
        p = ast.parse(inp).body[0]
        stack_size = process(consts, bytecode, p, stack_size)
        bytecode.append(RETURN_VALUE)
        bytecode = bytes(bytearray(bytecode))
        consts = tuple(consts)
        if PY3:
            code = types.CodeType(oldcode.co_argcount, oldcode.co_kwonlyargcount, oldcode.co_nlocals, stack_size, oldcode.co_flags, bytecode, consts, oldcode.co_names, oldcode.co_varnames, oldcode.co_filename, 'f', oldcode.co_firstlineno, b'')
            f.__code__ = code
        else:
            code = types.CodeType(oldcode.co_argcount, oldcode.co_nlocals, stack_size, oldcode.co_flags, bytecode, consts, oldcode.co_names, oldcode.co_varnames, oldcode.co_filename, 'f', oldcode.co_firstlineno, '')
            f.func_code = code
        return f
    

    This has the distinct advantage of generating essentially the same function as eval, and it scales almost exactly as well as compile+eval (the compile step is slightly slower than eval's, and eval will precompute anything it can (1+1+x gets compiled as 2+x).

    For comparison, eval finishes your 20k test in 0.0125 seconds, and makefunction finishes in 0.014 seconds. Increasing the number of iterations to 2,000,000, eval finishes in 1.23 seconds and makefunction finishes in 1.32 seconds.

    An interesting note, pypy recognizes that eval and makefunction produce essentially the same function, so the JIT warmup for the first accelerates the second.

提交回复
热议问题