«Оптимизируем» функции на уровне AST

Python предоставляет программисту огромное пространство свободы. Увы, обычно это довольно дорогая в плане производительности свобода, зато при правильном применении иногда она позволяет творить сущую магию. Но сегодня мы поговорим не о таких вот «богоугодных» применениях свободы, а о том, что никогда не стоит использовать в прикладном программировании — о модификациях кода на уровне AST.

Да кто такой этот ваш AST?

Обычно принято говорить, что CPython выполняет код построчно. Это не совсем правда, и прежде, чем идти дальше, я должен сделать небольшую паузу и кратко описать схему его работы.

Прежде всего, любой код — это последовательность символов в некоем файле (в случае REPL — последовательность, введенная в консоли). Как бы по-программистски она ни выглядела, она не является программой, то есть не может быть выполнена компьютером, пока с ней не будет проделан ряд шагов.

Первый шаг — парсинг текста. Если мы представим парсер как функцию, то на вход она должна получать «сырой» текст, а на выходе давать последовательность неких «токенов», то есть значащих элементов нашей программы. Токенами могут быть различные «захардкоженные» (for, if, True и т. д.) ключевые слова, названия сущностей (переменных, функций, классов), строковые литералы и прочее. Парсинг выполняется по некоторому набору правил, называемому абстрактной грамматикой. В случае CPython мы не имеем возможности, к сожалению, как-то влезть в процесс парсинга, скажем, чтобы добавить новый набор ключевых слов (хотя, поговаривают, лазейки всё же есть).

Второй шаг — построение на основе полученных из предыдущего шага токенов синтаксического дерева (или, как его принято называть у нас в деревне, AST). AST — это, как очевидно по названию, некая древовидная структура данных, где каждая нода получена из какого-нибудь токена, который мы извлекли на этапе парсинга. Эта нода может быть либо терминальной, либо содержащей ссылки на нижележащие ноды. Скажем, нода, представляющая блок условия, может содержать в себе список нод, описывающих действия программы в случае, если это условие выполняется. Ноды в AST являются типизированными, то есть они имеют разный набор свойств, когда представляют разные виды токенов.

Третий шаг — генерация на основе AST графа потока управления (Control Flow Graph). Очень грубо выражаясь, на этом шаге набор элементов исходного кода превращается в последовательность операций.

Четвертый и пятый шаги — компиляция графа в байт-код и его выполнение на виртуальной машине. Тут наши полномочия всё, код вступает в законную силу и изменению не подлежит (на самом деле подлежит, но надо упороться чуть посильнее).

Именно на втором шаге мы будем проводить свои «оптимизации» (я не случайно беру это слово в кавычки, причину объясню ниже). По сути, они сводятся к трем типам возможных действий: добавление, удаление или модификация уже существующих нод (скажем, подмена одного строкового литерала другим). Работа с синтаксическим деревом является, с одной стороны, достаточно простым, а с другой — очень мощным средством, позволяющим писать программы, которые буквально переписывают свой исходный код на лету (на самом деле нет, но выглядит как будто бы так).

Что очень важно понимать про AST: оно «глупое». Оно не знает, что на самом деле лежит в ваших переменных, это определяется во время выполнения. По сути всё, что оно из себя представляет, — это некая чуть более упорядоченная форма бытия того самого текстового файлика, что вы подали на вход программы.

Минздрав предупреждает

«не должен находиться у тебя [...] волшебник [...] ибо мерзок пред Господом всякий, делающий это, и за сии‐то мерзости Господь, Бог твой, изгоняет их от лица твоего» (Второзаконие 18:10-12)

Все эти наши штуки с AST интересные и вроде бы кажутся даже полезными, но не стоит поддаваться искушению применять такую магию в мире маглов. Здесь я опишу только общие соображения, почему это так, но дальше вы можете увидеть и другие проблемы с каждой из предложенных модификаций.

1. Фактически выполняться будет не тот код, что вы видите в файле с программой. У программистов пухнут головы уже от работы с тем кодом, что они видят, а тут, представьте, уже написанный код начинает себя вести не так, как выглядит.

2. Вы должны держать в голове нюансы реализации конкретной «оптимизации» и учитывать возможные побочные эффекты. Это вместо того, чтобы сосредоточиться на том, что должна делать ваша программа.

3. Средства интроспекции Python начинают врать. Допустим, мы добавили несколько нод в AST. Они не привязаны к конкретным строчкам исходного файла с кодом. А теперь вспомните, как выглядит, скажем, обратная трассировка, когда в коде происходят исключения. Там подробно указывается, в какой строчке кода и с какой операцией произошла ошибка. Если эта ошибка случится при выполнении кода из этих новых нод, трассировка вас просто обманет.

4. Реальная выгода от большинства модификаций обычно крайне мала. Да, вам может показаться хорошей идеей применить такие модификации, как, скажем, раскрытие циклов или инлайнинг функций. Но увы, как показала практика, реальное ускорение кода при этом обычно находится на грани погрешности измерений. Если вам не хватает скорости в программах, написанных на Python, то обычно самое лучшее решение — переписать их на чем-то еще.

Общий порядок действий

Итак, мы уже поняли, что хотим внедриться во второй этап работы интерпретатора и подсунуть ему модифицированное AST. Но не тут-то было: оказывается, мы не можем сделать это напрямую. По факту мы для конкретного объекта функции вручную повторяем весь интерпретаторский порядок действий, начиная с получения текста исходного кода и заканчивая генерированием байт-кода.

Чтобы вам проще было ориентироваться в коде, здесь и далее мы всегда будем модифицировать функцию с именем kek, но ее содержимое будет меняться. Вот в общих чертах, что мы с ней делаем:

1) Получаем исходный код при помощи функции из модуля inspect.

import ast, inspect


source = inspect.getsource(kek)

2) Заново парсим исходный код, получая на выходе то самое AST.

tree = ast.parse(source)

Если вам любопытно, можете взглянуть, как выглядит дерево, запринтив его в консоль:

print(ast.dump(tree, indent=4))

3) Что-то делаем с полученным AST. На выходе у нас должно получиться другое AST, которое мы обозначим как 'new_tree'. По сути, это единственная часть в нашем уравнении, которая будет отличаться для разных модификаций, так что подробнее об этом поговорим ниже.

4) Полученное AST скармливаем встроенной функции compile(), чтобы получить байт-код.

bytecode = compile(new_tree, filename=inspect.getfile(kek), mode='exec')

5) Для получения нового объекта функции необходимо выполнить байт-код с инициализацией функции kek. Для этого используем встроенную функцию exec().

namespace = {}
exec(bytecode, namespace)
new_kek = namespace['kek']

Обратите внимание на переменную namespace. Если в определении функции kek фигурировали какие-то объекты, определенные в модуле, то этими объектами нужно заполнить словарь пространства имен. Но в нашем упрощенном примере кода мы оставим его пустым.

Учитывайте, что если в коде вашей функции встречается директива global, то она будет привязана к тому «виртуальному» модулю, в котором выполнялся код через функцию exec(). Изменяя какой-то глобальный объект через new_kek и какую-нибудь еще функцию вы можете в обнаружить, что объекты, с которыми они работают — разные. В общем, не используйте global.

6) Если всё пошло по плану, то вы получили новый объект функции с внесенными модификациями. Попробуйте выполнить его!

new_kek('lol')

Что же, пора приступить к модификациям.

Простейшая модификация — выключатель для ифов

В пакете ast есть удобный инструмент для модифицирования узлов синтаксического дерева — класс NodeTransformer. У него есть метод, который запускает обход дерева. Каждую ноду при этом можно либо удалить, либо заменить на что-то другое. Для использования методов мы должны отнаследоваться от NodeTransformer и переопределить там метод, модифицирующий конкретный тип нод — if. Созданный под наши цели класс будет выглядеть так:

systems = ('linux': True, 'windows': False)

class RewriteIfs(ast.NodeTransformer):
    def visit_If(self, node):
        try:
            variable_name = node.test.id
            if variable_name in systems:
                if not systems[variable_name]:
                    return None
                else:
                    return node.body
            return node
        except Exception:
            return node

Запустим обход:

new_tree = RewriteIfs().visit(tree)

После перекомпиляции мы получим функцию с измененным поведением. При обходе мы побывали во всех нодах условий, и если условие было унарным (то есть в нем присутствовал только один операнд), а также если название переменной операнда фигурировало в словаре systems, то с этим блоком if будет выполнено одно из двух действий:

1) Если в словаре названию данного операнда соответствует False, то весь блок будет уничтожен — для этого метод visit_If при посещении соответствующей ноды должен вернуть None.

2) Если в словаре по тому же ключу лежит True, то мы берем содержимое блока if и заменяем на него сам if. То есть код из блока if начинает выполняться без условий. Функция, которая выглядит вот так:

def kek():
  print('begin')
  if linux:
    print('its linux')
  if windows:
    print('its windows')
  print('end')

... после модификации превратится вот в это:

def new_kek():
  print('begin')
  print('its linux')
  print('end')

Практически применить эту штуку можно, к примеру, если мы не хотим постоянно проверять в рантайме некое условие, которое точно известно на момент начала выполнения программы и не должно измениться до конца. Это может быть что-то, извлеченное из конфига, или какой-то платформозависимый кусок кода. Но лучше для таких целей всё же использовать какую-нибудь фабрику.

Оптимизация хвостовой рекурсии

Хвостовая рекурсия — это такая рекурсия, где функция возвращает результат вызова самой себя непосредственно в блоке return. Многие компиляторы других языков программирования умеют автоматически оптимизировать хвостовую рекурсию, превращая ее алгоритм в итеративный. Но в CPython такой оптимизации нет by design, по причинам, которые перечислил сам Гвидо ван Россум. Спорить с Гвидо мы не будем, его причины отказа от оптимизации вполне весомые, но всё же развлечения ради напишем свой оптимизатор. К сожалению, у нас не получится сделать это так же бесшовно, как эта функциональность могла бы выглядеть, будучи встроенной непосредственно в интерпретатор.

«Исправлять» рекурсию мы будем следующим образом:

1. Помещаем рекурсивную функцию в другую функцию, внутри которой первая будет вызываться внутри бесконечного цикла.

2. Правим AST рекурсивной функции. Нам нужно заменить внутри нее вызовы функций с тем же названием на блок return, возвращающий аргументы, которые функция передавала сама себе.

3. При этом нам нужно как-то отличать ситуацию обычного возвращения переменных из функции от такой вот подмены. Поэтому перед возвращением мы проставляем специальный флаг, обозначающий, что вот это сейчас было совсем не обычное возвращение результата работы функции, а нечто, символизирующее рекурсивный вызов.

4. На уровне цикла из п. 1 мы на каждой итерации вызываем измененную функцию и смотрим на флаг. Если он «поднят», значит на самом деле функция хотела вызвать саму себя, но в результате исправления что-то вернула. Поэтому мы просто переходим на следующую итерацию цикла и там передаем в функцию то, что она нам вернула. Если же флаг «опущен», значит мы достигли «дна» рекурсии и это тот результат, который мы должны были получить. Возвращаем его. Мы восхитительны!

Теперь то же самое, но немного подробнее.

1) Пишем функцию, внутрь которой мы поместим оригинал. Выглядеть она будет примерно так:

def superfunction(*args, **kwargs):
    {{kek}}
    is_recursion = False
    while True:
        result = kek(*args, **kwargs)
        if not is_recursion:
            return result
        is_recursion = False
        args, kwargs = result

В этом примере мы буквально вставим текст исходной функции на место {{kek}}. Да, мы сейчас работаем с кодом как со строкой. Нет, так делать неправильно, и существуют более разумные варианты, как это можно сделать. Если есть идеи, как, то, может, расскажете в комментариях?

2, 3 и 4) Вставляем текст исходной функции в обертку из п. 1, после чего парсим результат и исправляем AST. Допустим, исходная функция выглядела так:

def kek(number):
    if number >= 10000000:
        return number
    return kek(number + 1)

Обычно при вызове этой функции поднималось бы исключение, связанное с превышением максимального уровня глубины рекурсии:

>>> kek(1)
...
RecursionError: maximum recursion depth exceeded in comparison

После «вклеивания» в обертку мы получим примерно следующее:

def superfunction(*args, **kwargs):
    is_recursion = False
    def kek(number):
        if number >= 10000000:
            return number
        return kek(number + 1)
    while True:
        result = kek(*args, **kwargs)
        if not is_recursion:
            return result
        is_recursion = False
        args, kwargs = result

Обратите внимание, мы это делаем чисто на уровне текста! На выходе мы получаем переменную, содержащую строку с кодом, показанным выше. Обозначим ее далее как text_of_new_code.

Парсим исправленный код и получаем объект AST:

temp_tree = ast.fix_missing_locations(ast.parse(text_of_new_code))

Теперь самое время написать класс-обходчик нод дерева, который будет уделять особое внимание нодам с блоками return:

class RewriteReturns(ast.NodeTransformer):
    def visit_Return(self, node):
        if self.is_recursion('kek', node):
            _nonlocal = ast.Nonlocal(names=['is_recursion'])
            flag = ast.Assign(targets=[ast.Name(id='is_recursion', ctx=ast.Store())], value=ast.Constant(value=True))
            new_return_node = self.get_new_return_node(node)
            return [_nonlocal, flag, new_return_node]
        return node

    def is_recursion(self, function_name, node):
        try:
            if node.value.func.id == function_name:
                return True
            return False
        except Exception:
            return False

    def get_new_return_node(self, node):
        args = node.value.args
        kwargs = node.value.keywords
        dict_keys = [ast.Constant(value=keyword.arg) for keyword in kwargs]
        dict_values = [keyword.value for keyword in kwargs]
        new_node = ast.Return(
            value=ast.Tuple(
                elts=[
                    ast.Tuple(elts=args, ctx=ast.Load()),
                    ast.Dict(
                        keys=dict_keys,
                        values=dict_values,
                    ),
                ], ctx=ast.Load(),
            ), ctx=ast.Load(),
        )
        return new_node

Запустим обход дерева:

new_tree = RewriteReturns().visit(temp_tree)

Если расшифровать полученное дерево, то оно будет представлять собой примерно такую функцию:

def superfunction(*args, **kwargs):
    is_recursion = False
    def kek(number):
        if number >= 10000000:
            return number
        nonlocal is_recursion
        is_recursion = True
        return ((number + 1, ), {})
    while True:
        result = kek(*args, **kwargs)
        if not is_recursion:
            return result
        is_recursion = False
        args, kwargs = result

После ее компилирования мы получим функцию, которая ведет себя как оригинал, только вдруг куда-то исчезают ограничения на глубину рекурсии:

>>> new_kek(1)
10000000

А с этим есть декораторы?

Поздним вечером и, вероятно, в не очень здравом состоянии рассудка я написал небольшую библиотечку, которая делает вышеописанные штуки с помощью декораторов. Сразу предупреждаю, она не предназначена для практического использования, написана криво и никак не протестирована. Но вроде бы работает.

Поставить можно так:

$ pip install astrologic

Управление if'ами серьезно упрощается, сравните сами. Эта функция:

from astrologic import switcher


@switcher(a=False, b=True)
def function():
  print('begin')
  if a:
    print('block a')
  if b:
    print('block b')
  print('end')

Преобразуется в:

def function():
  print('begin')
  print('block b')
  print('end')

Также можно всего одним декоратором навесить автоматическую оптимизацию хвостовой рекурсии. Вот так:

from astrologic import no_recursion
@no_recursion
def c(b):
    if b == 100000000:
        return b
    return c(b + 1)
print(c(1))
# Оно не падает.

Итак, теперь вы знаете еще несколько вещей, которые делать не надо. Надеюсь, вас это не остановит.