Inlining Python Function

Question

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.

Stellular

Thanks for the interesting question, but I am curious about your intended application for this transformation. – Counterstatement

It may make more sense to do this transformation at the bytecode level. (If you're at the point where you're putting this level of effort into microoptimizations, though, using something like Cython will probably give you more bang for your buck.) – Alo

Hmm.. this doesn't seem like the normal definition of inlining (i.e. alpha-conversion + replacing the call with the code). You seem to be doing some sort of eager evaluation..? – Jamisonjammal

@Counterstatement Specifically I'm writing a library whose goal is to perform dynamic syntactic alteration to support other code generation libraries, such as numba.jit and tangent. Hence why I need the result to be a valid Python function. If I can dynamically produce equivalent Python functions that don't use certain code features, then I don't need to rely on those libraries to support those code features (such as function calls). – Stellular

@scnerd: Numba currently doesn't support exception handling, but it does support function calls, so transforming function calls by using exception handling sounds unhelpful. – Alo

@Jamisonjammal I'm no expert in compilers, so you're very possibly right. I only used the term "inlining" because it looks very similar to the intuitive equivalent of #pragma inline in C. Could you provide some reference/article/tutorial regarding the ideas you mention (alpha-conversion, etc.) that would clarify the distinction between what I'm asking for and how actual code inlining works? – Stellular

@Alo Good point. Both function calls and exception handling are fairly dynamic, complicated operations for a code generation library to work with. Consider a hypothetical library like numba that additionally doesn't support function calls; any idea how to solve this question in that case? – Stellular

Tangent doesn't support mutable objects, but it also supports function calls, so this doesn't seem helpful for tangent either. – Alo

@scnerd: Probably the easiest thing to get correct is to only support return at the end of the function. I know you're trying to reduce limitations, but still. Probably the least invasive option if you want to support return in other positions is to use exception handling, but that has its own limitations if the called function uses with or try, or if one of the exceptions you use bubbles out of somewhere unexpected. – Alo

Aside from that, you could always do a super-duper-invasive structured program theorem-style rewriting where you desugar everything and convert the function into a giant while loop with a giant branching construct inside. – Alo

Normally inlining means to conceptually replace a call to function f with argument a -- i.e. f(a), where f is defined by def f(x): .. with E[x->a] :- code-of-f, meaning execute the code of f with f's x-variable bound to a. Alpha-conversion simply means that variables in f needs to be renamed so they don't conflict with existing variables in the new location. Wikipedia has a general example (en.wikipedia.org/wiki/Inline_expansion#Implementation). LISP macros would be a better model than C-type #defines (cs.cmu.edu/Groups/AI/html/cltl/clm/node97.html) – Jamisonjammal

Answer 1

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.

Reptilian

Oh, I like that approach. It'd require making sure that the exception type is available within the function body, but I can do that pretty easily. I'm going to give this a bit longer in case other people have suggestions, but I'd be happy to accept this answer. – Stellular

@scnerd: You're going to have problems with except blocks inside the inlined function. – Alo

@Alo True, if they are bare except: or except BaseException: blocks. I could inherit from BaseException to minimize the risk, but it's certainly not foolproof. Under proper coding practices, though, this exception approach is safer and more general than my original for/break approach – Stellular

Answer 2

The only reasonable way on source level I see, simplified:

• Parse the source into some AST (or just use the built-in AST).
• Copy a subtree representing the function's body.
• Rename the variables in the subtree, e.g. by adding an unique prefix.
• At the call site, replace all passed arguments with assignments using the function's new variable names.
• Remove the call and replace it with the function body you've prepared.
• Serialize the AST back to source.

What poses real problems:

• Generator functions; just don't inline them.
• Returns from under try/finally that need to run the finally part. Might be pretty hard to rewrite correctly; imho, best left in-unlined.
• Returns from under context managers that need to run the __exit__ parts. While not impossible, it's also tricky to rewrite preserving the semantics; likely also best left un-inlined.
• Mid-function returns, especially from within multiple loop constructs. You might need to replace them with an extra variable and thread it into every condition of every while statement, and likely to add a conditional break to for statements. Again, not impossible but likely best left un-inlined.

Serajevo

@thebjorn: Indeed! Though tail recursion can be replaced with a loop :) – Serajevo

Sure, and you could translate it all to CPS (continuation-passing-style) ;-) -- but whatever you try you'll very quickly run into the fact that Python is not friendly to simple forms of inlining. The PyPy inliner doesn't even try to do inlining at the source level, and those guys have worked on this for a while. – Jamisonjammal

@thebjorn: Afaict Python here is an intermediate form: «Specifically I'm writing a library whose goal is to perform dynamic syntactic alteration to support other code generation libraries». I think simple forms of inlining (e.g. inlining a a function with a single return as the last statement) is doable. With something more advanced, complexity (and the chance of bugs) snowballs. – Serajevo

For recursive functions, so far I've just been tracking the recursion depth, so you actually have _g_0, _g_1, etc., and the recursion limit can be explicitly capped (and should be a really low cap). Pretty sure generator functions are impossible to inline syntactically, so I've taken a "you better know what you're doing" approach and treated a call to a generator function f(x) to be equivalent to list(f(x)), which makes it inlinable. – Stellular

In cases of finally and context manager __exit__'s, they should get run properly as the return-exception cascades up past them, shouldn't they? That's one of the big points of those two paradigms is to be resilient to being exception'ed past, right? – Stellular

@scnerd: Exactly; they run as if the returned value is stored, then the post-action is run, and then the stored value is returned. This is why they ^ are tricky to inline. – Serajevo