Search code examples
pythoninlineabstract-syntax-treecompiler-optimization

Inlining Python Function


In a C program, inlining a function is a fairly intuitive optimization. If the inlined function's body is sufficiently small, you end up saving the jump to the function and creation of the stack frame, and you store the return value wherever the function's result would have been stored, jumping to the end of the inlined function's "body" rather than long-jumping to the return pointer.

I'm interested in doing the same thing in Python, converting two python functions into another valid python function where the first got "inlined" into the second. An ideal solution to this might look something like the following:

def g(x):
    return x ** 2

def f(y):
    return g(y + 3)

# ... Becomes ...

def inlined_f(y):
    return (y + 3) ** 2

Clearly, in a language as dynamic as Python, this isn't trivial to do automatically. The best generic solution I have come up with is to use dict to capture the arguments passed to the function, wrap the function body in a one-iteration for loop, use break to jump to the end of the function, and replace uses of arguments with indexes into the argument dictionary. The result looks something like the following:

def inlined_f(y):
    _g = dict(x=y + 3)
    for ____ in [None]:
        _g['return'] = _g['x'] ** 2
        break
    _g_return = _g.get('return', None)
    del _g
    return _g_return

I don't care that it's ugly, but I do care that it doesn't support returns from within loops. E.g.:

def g(x):
    for i in range(x + 1):
        if i == x:
            return i ** 2
    print("Woops, you shouldn't get here")

def inlined_f(y):
    _g = dict(x=y + 3)
    for ____ in [None]:
        for _g['i'] in range(_g['x'] + 1):
            if _g['i'] == _g['x']:
                _g['return'] _g['i'] ** 2
                break  # <-- Doesn't exit function, just innermost loop
        print("Woops, you shouldn't get here")
    _g_return = _g.get('return', None)
    del _g
    return _g_return

What approach could I take to this problem that avoids needing to use break to "jump" out of the inlined function's body? I'd also be open to an overall better, generic approach could I take to inline one Python function into another.

For reference, I'm working at the AST (abstract syntax tree) level, so using parsed Python code; clearly, outside of literal values, I don't know what value or type anything will have while performing this transformation. The resulting inlined function must behave identically to the original functions, and must support all features typically available when calling a function. Is this even possible in Python?


EDIT: I should clarify since I used the tag "optimization", that I'm not actually interested in a performance boost. The resulting code does not need to be faster, it just must not call the inlined function while still behaving identically. You can assume that both functions' source code is available as valid Python.


Solution

  • Probably the closest analog to a return would be raising an Exception, which would work to pop out of nested loops to the top of the "inlined function".

    class ReturnException(Exception):
        pass
    
    
    g = dict(x=y + 3)
    try:
        for j in some_loop:
            for _g['i'] in range(_g['x'] + 1):
                if _g['i'] == _g['x']:
                    raise ReturnException(_g['i'] ** 2)
    except ReturnException as e:
        _g['return'] = e.message
    else:
        _g['return'] = None
    

    I don't know how much overhead is associated with exceptions though or if that would be faster than simply calling the function.