Путь хакера — решение задачи с codechef на питон. С машинным кодом

Материал из DISCOPAL
Перейти к: навигация, поиск
 
(не показано 7 промежуточных версий 2 участников)
Строка 1: Строка 1:
 
В процессе написания [[../Python-оптимизация алгоритма динамического программирования из codechef]], когда удалось получить питон-решение проходящее таймлимиты на PyPy, возник вопрос-челлендж, можно ли все-таки пройти этот квест на обычном питоне («CPython»).
 
В процессе написания [[../Python-оптимизация алгоритма динамического программирования из codechef]], когда удалось получить питон-решение проходящее таймлимиты на PyPy, возник вопрос-челлендж, можно ли все-таки пройти этот квест на обычном питоне («CPython»).
  
Казалось, что это невозможно, ведь CPython грубо получает таймлимит больше в 2.5 раза, чем для PyPy (но это совсем небольшая фора), и для PyPy мы еле-еле вписались, а по нашим оценкам, наш же алгоритм на CPython работал на порядок медленней. Возможно ожидалась какая-то магия с использованием модулей <tt>numpy</tt> или <tt>collections</tt>, чтобы по сути, дергать питоном написанные на си или даже высокообогащенном оружейном фортране функции.
+
Казалось, что это невозможно, ведь CPython грубо получает таймлимит больше в 2.5 раза, чем для PyPy (но это совсем небольшая фора), и для PyPy мы еле-еле вписались, а по нашим оценкам, наш же алгоритм на CPython работал на порядок медленней. Возможно ожидалась какая-то магия с использованием модулей <tt>numpy</tt> или <tt>collections</tt>, чтобы по сути, дергать питоном функции, написанные на си или даже высокообогащенном оружейном фортране.
  
 
Другие методы оптимизации — отдельные скомпилированные модули на С, компиляция Cython или Nuitka, тут явно невыполнимы — на вход принимается только один питон-файл.
 
Другие методы оптимизации — отдельные скомпилированные модули на С, компиляция Cython или Nuitka, тут явно невыполнимы — на вход принимается только один питон-файл.
  
Но оказалось, есть еще и путь хакера.
+
Но оказалось, есть еще и [https://www.codechef.com/viewsolution/64117884 путь хакера].
  
 
<source lang="python">
 
<source lang="python">
Строка 25: Строка 25:
  
  
func_bytes = b’AWA\x89\xf7AVA\x89\xfeAUA\x89\xd5ATUSH\x83\xec8\x85\xd2u5Hc\xc7E1\xe4\x8b,\x81\x85\xedy-C\x8d\x04wA' \
+
func_bytes = (b'AWA\x89\xf7AVA\x89\xfeAUA\x89\xd5ATUSH\x83\xec8\x85\xd2u5Hc\xc7E1'
            b'\x8dDE\x00H\x98M\x89$\xc1H\x83\xc48L\x89\xe0[]A\\A]A^A_\xc3\x0f\x1f\x80\x00\x00\x00\x00\xbd\t\x00\x00' \
+
              b'\xe4\x8b,\x81\x85\xedy-C\x8d\x04wA\x8dDE\x00H\x98M\x89$\xc1H\x83\xc48L'
            b'\x00A\x8d~\x01D\x89\xc0H\x8bt$p)\xf8E\x85\xffH\x98A\x0f\x95\xc2L\x8d\x1c\xc6A9\xf8tk\x8d\x04?1\xdb' \
+
              b'\x89\xe0[]A\\A]A^A_\xc3\x0f\x1f\x80\x00\x00\x00\x00\xbd\t\x00\x00\x00A\x8d~'
            b'E1\xe4\x89D$\x08\x83\xfb\x02t\x19E\x84\xd2u\x14\x83\xfb\x07\x0f\x87|\x00\x00\x00\xb8\xa8\x00\x00\x00H' \
+
              b'\x01D\x89\xc0H\x8bt$p)\xf8E\x85\xffH\x98A\x0f\x95\xc2L\x8d\x1c\xc6A9\xf8t'
            b'\x0f\xa3\xd8sq9\xddu51\xd2\xbe\x01\x00\x00\x00E\x85\xedu)\x8bD$\x08\x01\xf0\x8d\x04BH\x98I\x8b\x04' \
+
              b'k\x8d\x04?1\xdbE1\xe4\x89D$\x08\x83\xfb\x02t\x19E\x84\xd2u\x14\x83'
            b'\xc1H\x83\xf8\xfftdI\x01\xc4\x83\xc3\x019\xdd}\xb3\xe9U\xff\xff\xff\x0f\x1fD\x00\x00M\x03#\xeb\xea' \
+
              b'\xfb\x07\x0f\x87|\x00\x00\x00\xb8\xa8\x00\x00\x00H\x0f\xa3\xd8sq9\xddu51'
            b'\x0f\x1f\x001\xc0E1\xe4\xba\xa8\x00\x00\x00\x83\xf8\x02t\x10E\x84\xd2u\x0b\x83\xf8\x07w\nH\x0f\xa3' \
+
              b'\xd2\xbe\x01\x00\x00\x00E\x85\xedu)\x8bD$\x08\x01\xf0\x8d\x04BH\x98I\x8b'
            b'\xc2s\x04I\x83\xc4\x01\x83\xc0\x019\xc5}\xe0\xe9\x19\xff\xff\xff\x909\xdd\x0f\x95\xc2E\x85\xed\x0f' \
+
              b'\x04\xc1H\x83\xf8\xfftdI\x01\xc4\x83\xc3\x019\xdd}\xb3\xe9U\xff\xff\xff\x0f'
            b'\x95\xc01\xf6\t\xc2\x0f\xb6\xd2\xeb\x8b\x0f\x1f@\x00D\x88T$/H\x83\xec\x08L\x89\\$(\xfft$xL\x89L$(D' \
+
              b'\x1fD\x00\x00M\x03#\xeb\xea\x0f\x1f\x001\xc0E1\xe4\xba\xa8\x00'
            b'\x89D$8H\x89L$ \x89|$\x1c\xe8\xae\xfe\xff\xffI\x01\xc4XZ\x8b|$\x0cH\x8bL$\x10D\x8bD$(L\x8bL$\x18L' \
+
              b'\x00\x00\x83\xf8\x02t\x10E\x84\xd2u\x0b\x83\xf8\x07w\nH\x0f\xa3\xc2s\x04I'
            b'\x8b\\$ D\x0f\xb6T$/\xe9M\xff\xff\xff'
+
              b'\x83\xc4\x01\x83\xc0\x019\xc5}\xe0\xe9\x19\xff\xff\xff\x909\xdd\x0f\x95'
 +
              b'\xc2E\x85\xed\x0f\x95\xc01\xf6\t\xc2\x0f\xb6\xd2\xeb\x8b\x0f\x1f@\x00D\x88T$'
 +
              b'/H\x83\xec\x08L\x89\\$(\xfft$xL\x89L$(D\x89D$8H\x89L$ \x89|$\x1c\xe8\xae\xfe'
 +
              b'\xff\xffI\x01\xc4XZ\x8b|$\x0cH\x8bL$\x10D\x8bD$(L\x8bL$\x18L\x8b\\$ D'
 +
              b'\x0f\xb6T$/\xe9M\xff\xff\xff')
  
 
assert len(func_bytes) <= mmap.PAGESIZE
 
assert len(func_bytes) <= mmap.PAGESIZE
Строка 63: Строка 67:
 
Идея — будем действовать как JIT. Мы можем переписать кусок кода на си, взять получившийся машинный код, внутри питона его записать в динамически выделенную память и запустить.
 
Идея — будем действовать как JIT. Мы можем переписать кусок кода на си, взять получившийся машинный код, внутри питона его записать в динамически выделенную память и запустить.
  
Мы можем выделить память с разрешением на чтение, запись и исполнение. Запись понятно для чего, а разрешение на исполнение означает что процессор может сделать jmp или call в эту область и не крешнуться (разрешение на чтение казалось бы не нужно, но как ни странно, на x86-64 для исполнения также нужно разрешение на чтение):
+
Мы можем выделить память с разрешением на чтение, запись и исполнение. Запись понятно для чего, а разрешение на исполнение означает, что процессор может сделать <tt>jmp</tt> или <tt>call</tt> в эту область и не крешнуться (разрешение на чтение казалось бы не нужно, но как ни странно, на x86-64 для исполнения также нужно разрешение на чтение):
  
 
<source lang="python">
 
<source lang="python">
Строка 105: Строка 109:
 
Теперь о том как получить нужный машинный код и как вообще выбрать что компилировать.
 
Теперь о том как получить нужный машинный код и как вообще выбрать что компилировать.
  
Наши возможности в том что мы можем запустить довольно ограничены — мы не можем запустить что-либо что использует какие-либо библиотеки, потому что мы загружаем не честный ELF, в котором можно подгружать библиотеки, делать релокации и многое другое, а только набор машинных инструкций. Это значит что мы не можем использовать ничего из libc — то есть без ввода/вывода, динамического выделения памяти, ничего из стандартной библиотеки плюсов — то есть без каких-либо структур данных сложнее статического массива.
+
Наши возможности в том что мы можем запустить довольно ограничены — мы не можем запустить что-либо что использует какие-либо библиотеки, потому что мы загружаем не «честный ELF», в котором можно подгружать библиотеки, делать релокации и многое другое, а только набор машинных инструкций. Это значит что мы не можем использовать ничего из <tt>libc</tt> — то есть без ввода/вывода, динамического выделения памяти, ничего из стандартной библиотеки плюсов — то есть без каких-либо структур данных сложнее статического массива.
  
 
Поэтому идеальный кандидат — функция calculate из оригинального кода на питоне. Она как раз использует только массивы фиксированного размера и простые операции:
 
Поэтому идеальный кандидат — функция calculate из оригинального кода на питоне. Она как раз использует только массивы фиксированного размера и простые операции:
Строка 202: Строка 206:
 
(глобальные переменные мы создаем на стороне питона и передаем как аргументы, также мы преобразовали многомерный массив dp в одномерный, чтобы не париться с указателями)
 
(глобальные переменные мы создаем на стороне питона и передаем как аргументы, также мы преобразовали многомерный массив dp в одномерный, чтобы не париться с указателями)
  
Дальше нам нужно это скомпилировать, и нужно передать некоторые нестандарные флаги. Обычно исполнение ELF начинается не с main, а с функции _start, которая вызывает функцию __libc_start_main из libc, которая подгружает переменные окружения, argv и многое другое, а только потом запускает main. Нам ничего такого не нужно, поэтому мы это отключаем с помощью флага -nostartfiles, и к тому же без этого компилятор ожидал бы от нас функцию main, которую мы писать не хотим. Также можем добавить -nolibc чтобы не было соблазна вызывать функции оттуда, хотя на результат это не повлияет. И конечно -O3 чтобы получить все оптимизации что есть.
+
Дальше нам нужно это скомпилировать, и нужно передать некоторые нестандарные флаги. Обычно исполнение ELF начинается не с <tt>main</tt>, а с функции <tt>_start</tt>, которая вызывает функцию <tt>__libc_start_main</tt> из <tt>libc</tt>, которая подгружает переменные окружения, <tt>argv</tt> и многое другое, а только потом запускает <tt>main</tt>. Нам ничего такого не нужно, поэтому мы это отключаем с помощью флага <tt>-nostartfiles</tt>, и к тому же без этого компилятор ожидал бы от нас функцию <tt>main</tt>, которую мы писать не хотим. Также можем добавить <tt>-nolibc</tt> чтобы не было соблазна вызывать функции оттуда, хотя на результат это не повлияет. И конечно <tt>-O3</tt> чтобы получить все оптимизации что есть.
  
 
<source lang="bash">
 
<source lang="bash">
Строка 208: Строка 212:
 
</source>
 
</source>
  
Можем проверить что в секции .text у нас лежит одна единственная функция (-d = вывод дизассемблера секции с кодом, -M intel = чтобы получить вывод в намного более приятном интеловском синтаксисе, а не в AT&T):
+
Можем проверить что в секции <tt>.text</tt> у нас лежит одна единственная функция (<tt>-d</tt> — вывод дизассемблера секции с кодом, <tt>-M intel</tt> — чтобы получить вывод в намного более приятном интеловском синтаксисе, а не в AT&T):
  
 
<source lang="html4strict">
 
<source lang="html4strict">
Строка 242: Строка 246:
 
</source>
 
</source>
  
Также стоит заметить что наш код не использует никаких глобальных переменных, а значит нам нужна только секция .text из ELF. И еще все инструкции jmp и call используют относительные оффсеты для прыжков, поэтому наш код может быть загружен по любому адресу памяти, что нам и нужно, так как мы не можем полностью контроллировать какой адрес вернет mmap.
+
Также стоит заметить что наш код не использует никаких глобальных переменных, а значит нам нужна только секция <tt>.text</tt> из ELF. И еще все инструкции <tt>jmp</tt> и <tt>call</tt> используют относительные оффсеты для прыжков, поэтому наш код может быть загружен по любому адресу памяти, что нам и нужно, так как мы не можем полностью контроллировать какой адрес вернет mmap.
  
Нам осталось сдампить байты нашей функции в отдельный файл (-O binary = вывести сырые байты, -j .text = работать с секцией .text):
+
Нам осталось сдампить байты нашей функции в отдельный файл (<tt>-O binary</tt> = вывести сырые байты, <tt>-j .text</tt> = работать с секцией <tt>.text</tt>):
  
 
<source lang="bash">
 
<source lang="bash">
Строка 253: Строка 257:
  
 
<source lang="python">
 
<source lang="python">
# with open('inline_bytes', 'rb') as f:
+
from pprint import pprint
# print(f.read())
+
  
func_bytes = b’AWA\x89\xf7AVA\x89\xfeAUA\x89\xd5ATUSH\x83\xec8\x85\xd2u5Hc\xc7E1\xe4\x8b,\x81\x85\xedy-C\x8d\x04wA' \
+
with open('inline_bytes', 'rb') as f:
            b'\x8dDE\x00H\x98M\x89$\xc1H\x83\xc48L\x89\xe0[]A\\A]A^A_\xc3\x0f\x1f\x80\x00\x00\x00\x00\xbd\t\x00\x00' \
+
    pprint(f.read())
            b'\x00A\x8d~\x01D\x89\xc0H\x8bt$p)\xf8E\x85\xffH\x98A\x0f\x95\xc2L\x8d\x1c\xc6A9\xf8tk\x8d\x04?1\xdb' \
+
            b'E1\xe4\x89D$\x08\x83\xfb\x02t\x19E\x84\xd2u\x14\x83\xfb\x07\x0f\x87|\x00\x00\x00\xb8\xa8\x00\x00\x00H' \
+
            b'\x0f\xa3\xd8sq9\xddu51\xd2\xbe\x01\x00\x00\x00E\x85\xedu)\x8bD$\x08\x01\xf0\x8d\x04BH\x98I\x8b\x04' \
+
            b'\xc1H\x83\xf8\xfftdI\x01\xc4\x83\xc3\x019\xdd}\xb3\xe9U\xff\xff\xff\x0f\x1fD\x00\x00M\x03#\xeb\xea' \
+
            b'\x0f\x1f\x001\xc0E1\xe4\xba\xa8\x00\x00\x00\x83\xf8\x02t\x10E\x84\xd2u\x0b\x83\xf8\x07w\nH\x0f\xa3' \
+
            b'\xc2s\x04I\x83\xc4\x01\x83\xc0\x019\xc5}\xe0\xe9\x19\xff\xff\xff\x909\xdd\x0f\x95\xc2E\x85\xed\x0f' \
+
            b'\x95\xc01\xf6\t\xc2\x0f\xb6\xd2\xeb\x8b\x0f\x1f@\x00D\x88T$/H\x83\xec\x08L\x89\\$(\xfft$xL\x89L$(D' \
+
            b'\x89D$8H\x89L$ \x89|$\x1c\xe8\xae\xfe\xff\xffI\x01\xc4XZ\x8b|$\x0cH\x8bL$\x10D\x8bD$(L\x8bL$\x18L' \
+
            b'\x8b\\$ D\x0f\xb6T$/\xe9M\xff\xff\xff'
+
 
</source>
 
</source>
 +
 +
Какой будет вывод:
 +
 +
<source lang="bash">
 +
(b'AWA\x89\xf7AVA\x89\xfeAUA\x89\xd5ATUSH\x83\xec8\x85\xd2u5Hc\xc7E1'
 +
b'\xe4\x8b,\x81\x85\xedy-C\x8d\x04wA\x8dDE\x00H\x98M\x89$\xc1H\x83\xc48L'
 +
b'\x89\xe0[]A\\A]A^A_\xc3\x0f\x1f\x80\x00\x00\x00\x00\xbd\t\x00\x00\x00A\x8d~'
 +
b'\x01D\x89\xc0H\x8bt$p)\xf8E\x85\xffH\x98A\x0f\x95\xc2L\x8d\x1c\xc6A9\xf8t'
 +
b'k\x8d\x04?1\xdbE1\xe4\x89D$\x08\x83\xfb\x02t\x19E\x84\xd2u\x14\x83'
 +
b'\xfb\x07\x0f\x87|\x00\x00\x00\xb8\xa8\x00\x00\x00H\x0f\xa3\xd8sq9\xddu51'
 +
b'\xd2\xbe\x01\x00\x00\x00E\x85\xedu)\x8bD$\x08\x01\xf0\x8d\x04BH\x98I\x8b'
 +
b'\x04\xc1H\x83\xf8\xfftdI\x01\xc4\x83\xc3\x019\xdd}\xb3\xe9U\xff\xff\xff\x0f'
 +
b'\x1fD\x00\x00M\x03#\xeb\xea\x0f\x1f\x001\xc0E1\xe4\xba\xa8\x00'
 +
b'\x00\x00\x83\xf8\x02t\x10E\x84\xd2u\x0b\x83\xf8\x07w\nH\x0f\xa3\xc2s\x04I'
 +
b'\x83\xc4\x01\x83\xc0\x019\xc5}\xe0\xe9\x19\xff\xff\xff\x909\xdd\x0f\x95'
 +
b'\xc2E\x85\xed\x0f\x95\xc01\xf6\t\xc2\x0f\xb6\xd2\xeb\x8b\x0f\x1f@\x00D\x88T$'
 +
b'/H\x83\xec\x08L\x89\\$(\xfft$xL\x89L$(D\x89D$8H\x89L$ \x89|$\x1c\xe8\xae\xfe'
 +
b'\xff\xffI\x01\xc4XZ\x8b|$\x0cH\x8bL$\x10D\x8bD$(L\x8bL$\x18L\x8b\\$ D'
 +
b'\x0f\xb6T$/\xe9M\xff\xff\xff')
 +
</source>
 +
 +
И это нужно скопировать, чтобы можно было в codechef послать только один файл:
 +
 +
<source lang="python">
 +
func_bytes = (b'AWA\x89\xf7AVA\x89\xfeAUA\x89\xd5ATUSH\x83\xec8\x85\xd2u5Hc\xc7E1'
 +
              b'\xe4\x8b,\x81\x85\xedy-C\x8d\x04wA\x8dDE\x00H\x98M\x89$\xc1H\x83\xc48L'
 +
              b'\x89\xe0[]A\\A]A^A_\xc3\x0f\x1f\x80\x00\x00\x00\x00\xbd\t\x00\x00\x00A\x8d~'
 +
              b'\x01D\x89\xc0H\x8bt$p)\xf8E\x85\xffH\x98A\x0f\x95\xc2L\x8d\x1c\xc6A9\xf8t'
 +
              b'k\x8d\x04?1\xdbE1\xe4\x89D$\x08\x83\xfb\x02t\x19E\x84\xd2u\x14\x83'
 +
              b'\xfb\x07\x0f\x87|\x00\x00\x00\xb8\xa8\x00\x00\x00H\x0f\xa3\xd8sq9\xddu51'
 +
              b'\xd2\xbe\x01\x00\x00\x00E\x85\xedu)\x8bD$\x08\x01\xf0\x8d\x04BH\x98I\x8b'
 +
              b'\x04\xc1H\x83\xf8\xfftdI\x01\xc4\x83\xc3\x019\xdd}\xb3\xe9U\xff\xff\xff\x0f'
 +
              b'\x1fD\x00\x00M\x03#\xeb\xea\x0f\x1f\x001\xc0E1\xe4\xba\xa8\x00'
 +
              b'\x00\x00\x83\xf8\x02t\x10E\x84\xd2u\x0b\x83\xf8\x07w\nH\x0f\xa3\xc2s\x04I'
 +
              b'\x83\xc4\x01\x83\xc0\x019\xc5}\xe0\xe9\x19\xff\xff\xff\x909\xdd\x0f\x95'
 +
              b'\xc2E\x85\xed\x0f\x95\xc01\xf6\t\xc2\x0f\xb6\xd2\xeb\x8b\x0f\x1f@\x00D\x88T$'
 +
              b'/H\x83\xec\x08L\x89\\$(\xfft$xL\x89L$(D\x89D$8H\x89L$ \x89|$\x1c\xe8\xae\xfe'
 +
              b'\xff\xffI\x01\xc4XZ\x8b|$\x0cH\x8bL$\x10D\x8bD$(L\x8bL$\x18L\x8b\\$ D'
 +
              b'\x0f\xb6T$/\xe9M\xff\xff\xff')
 +
</source>
 +
 +
Вот и все!
 +
{{wl-publish: 2022-05-19 02:27:34 +0000 | StasFomin }}

Текущая версия на 02:27, 19 мая 2022

В процессе написания Blog:Advanced Algorithms/Python-оптимизация алгоритма динамического программирования из codechef, когда удалось получить питон-решение проходящее таймлимиты на PyPy, возник вопрос-челлендж, можно ли все-таки пройти этот квест на обычном питоне («CPython»).

Казалось, что это невозможно, ведь CPython грубо получает таймлимит больше в 2.5 раза, чем для PyPy (но это совсем небольшая фора), и для PyPy мы еле-еле вписались, а по нашим оценкам, наш же алгоритм на CPython работал на порядок медленней. Возможно ожидалась какая-то магия с использованием модулей numpy или collections, чтобы по сути, дергать питоном функции, написанные на си или даже высокообогащенном оружейном фортране.

Другие методы оптимизации — отдельные скомпилированные модули на С, компиляция Cython или Nuitka, тут явно невыполнимы — на вход принимается только один питон-файл.

Но оказалось, есть еще и путь хакера.

import copy
import ctypes
import mmap
import sys
 
DIGITS = 19
DP_SIZE = DIGITS * 2 * 2
 
 
def get_input():
    lines = sys.stdin.readlines()
    data = ' '.join(lines).strip()
 
    for token in data.split():
        yield token
 
 
func_bytes = (b'AWA\x89\xf7AVA\x89\xfeAUA\x89\xd5ATUSH\x83\xec8\x85\xd2u5Hc\xc7E1'
              b'\xe4\x8b,\x81\x85\xedy-C\x8d\x04wA\x8dDE\x00H\x98M\x89$\xc1H\x83\xc48L'
              b'\x89\xe0[]A\\A]A^A_\xc3\x0f\x1f\x80\x00\x00\x00\x00\xbd\t\x00\x00\x00A\x8d~'
              b'\x01D\x89\xc0H\x8bt$p)\xf8E\x85\xffH\x98A\x0f\x95\xc2L\x8d\x1c\xc6A9\xf8t'
              b'k\x8d\x04?1\xdbE1\xe4\x89D$\x08\x83\xfb\x02t\x19E\x84\xd2u\x14\x83'
              b'\xfb\x07\x0f\x87|\x00\x00\x00\xb8\xa8\x00\x00\x00H\x0f\xa3\xd8sq9\xddu51'
              b'\xd2\xbe\x01\x00\x00\x00E\x85\xedu)\x8bD$\x08\x01\xf0\x8d\x04BH\x98I\x8b'
              b'\x04\xc1H\x83\xf8\xfftdI\x01\xc4\x83\xc3\x019\xdd}\xb3\xe9U\xff\xff\xff\x0f'
              b'\x1fD\x00\x00M\x03#\xeb\xea\x0f\x1f\x001\xc0E1\xe4\xba\xa8\x00'
              b'\x00\x00\x83\xf8\x02t\x10E\x84\xd2u\x0b\x83\xf8\x07w\nH\x0f\xa3\xc2s\x04I'
              b'\x83\xc4\x01\x83\xc0\x019\xc5}\xe0\xe9\x19\xff\xff\xff\x909\xdd\x0f\x95'
              b'\xc2E\x85\xed\x0f\x95\xc01\xf6\t\xc2\x0f\xb6\xd2\xeb\x8b\x0f\x1f@\x00D\x88T$'
              b'/H\x83\xec\x08L\x89\\$(\xfft$xL\x89L$(D\x89D$8H\x89L$ \x89|$\x1c\xe8\xae\xfe'
              b'\xff\xffI\x01\xc4XZ\x8b|$\x0cH\x8bL$\x10D\x8bD$(L\x8bL$\x18L\x8b\\$ D'
              b'\x0f\xb6T$/\xe9M\xff\xff\xff')
 
assert len(func_bytes) <= mmap.PAGESIZE
 
func_buffer = mmap.mmap(-1, mmap.PAGESIZE, prot=mmap.PROT_READ | mmap.PROT_WRITE | mmap.PROT_EXEC)
 
func_type = ctypes.CFUNCTYPE(ctypes.c_longlong, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_int),
                             ctypes.c_int, ctypes.POINTER(ctypes.c_longlong), ctypes.POINTER(ctypes.c_longlong))
func_pointer = ctypes.c_void_p.from_buffer(func_buffer)
 
func = func_type(ctypes.addressof(func_pointer))
func_buffer.write(func_bytes)
 
token_generator = get_input()
next(token_generator)
 
dp = (ctypes.c_longlong * DP_SIZE)(*(-1 for _ in range(DP_SIZE)))
p10 = (ctypes.c_longlong * DIGITS)(*(10 ** i for i in range(DIGITS)))
 
for input_str in token_generator:
    N = len(input_str)
    input_digits = (ctypes.c_int * N)(*(int(ch) for ch in input_str))
 
    print(func(0, 0, 0, input_digits, N, copy.copy(dp), p10))

Идея — будем действовать как JIT. Мы можем переписать кусок кода на си, взять получившийся машинный код, внутри питона его записать в динамически выделенную память и запустить.

Мы можем выделить память с разрешением на чтение, запись и исполнение. Запись понятно для чего, а разрешение на исполнение означает, что процессор может сделать jmp или call в эту область и не крешнуться (разрешение на чтение казалось бы не нужно, но как ни странно, на x86-64 для исполнения также нужно разрешение на чтение):

func_buffer = mmap.mmap(-1, mmap.PAGESIZE, prot=mmap.PROT_READ | mmap.PROT_WRITE | mmap.PROT_EXEC)

Дальше мы пользуемся библиотекой ctypes в питоне — она позволяет из питона вызывать сишные функции. Для этого нам нужно указать какие аргументы наша функция будет принимать и что возвращать (чтобы питоновский код знал как нужно передавать аргументы перед тем как сделать call в наш код, а также как распарсить ответ):

func_type = ctypes.CFUNCTYPE(ctypes.c_longlong, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_int),
                             ctypes.c_int, ctypes.POINTER(ctypes.c_longlong), ctypes.POINTER(ctypes.c_longlong))

Вот как это соответствует нашей сишной функции:

int64_t calculate(int i, int taken, int ls, int *input_digits, int N, int64_t *dp, int64_t *p10)

Дальше нам нужно создать объект сишной функции и записать байты (о них дальше) в выделенную память:

func_pointer = ctypes.c_void_p.from_buffer(func_buffer)
func = func_type(ctypes.addressof(func_pointer))
func_buffer.write(func_bytes)

И после всего этого мы можем вызывать эту функцию как обычно, только нужно передавать аргументы нужных типов:

dp = (ctypes.c_longlong * DP_SIZE)(*(-1 for _ in range(DP_SIZE)))
p10 = (ctypes.c_longlong * DIGITS)(*(10 ** i for i in range(DIGITS)))
 
for input_str in token_generator:
    N = len(input_str)
    input_digits = (ctypes.c_int * N)(*(int(ch) for ch in input_str))
 
    print(func(0, 0, 0, input_digits, N, copy.copy(dp), p10))

Теперь о том как получить нужный машинный код и как вообще выбрать что компилировать.

Наши возможности в том что мы можем запустить довольно ограничены — мы не можем запустить что-либо что использует какие-либо библиотеки, потому что мы загружаем не «честный ELF», в котором можно подгружать библиотеки, делать релокации и многое другое, а только набор машинных инструкций. Это значит что мы не можем использовать ничего из libc — то есть без ввода/вывода, динамического выделения памяти, ничего из стандартной библиотеки плюсов — то есть без каких-либо структур данных сложнее статического массива.

Поэтому идеальный кандидат — функция calculate из оригинального кода на питоне. Она как раз использует только массивы фиксированного размера и простые операции:

def calculate(i, taken, ls):
    """
    i — индекс цифры
    taken — встречали ли простые цифры раньше
    ls — в текущем индексе можно делать полный перебор
    """
    global input_digits
    global dp
    global N
 
    limit = 9
 
    if not ls:
        limit = input_digits[i]
 
    res = 0
 
    for x in range(limit + 1):
        newtaken = taken or (x == 2) or (x == 3) or (x == 5) or (x == 7)
        newi = i + 1
 
        if newi == N:
            ret = newtaken
        else:
            newls = ls or (x != limit)
 
            if newls and newtaken:
                ret = p10[N - newi]
            else:
                cached = dp[newtaken][newls][newi]
 
                if cached != -1:
                    ret = cached
                else:
                    ret = calculate(newi, newtaken, newls)
 
        res += ret
 
    dp[taken][ls][i] = res
 
    return res

API ctypes позволяет нам запускать сишные функции — поэтому ее и напишем:

// file: inline.c
# include <stdint.h>
 
int64_t calculate(int i, int taken, int ls, int *input_digits, int N, int64_t *dp, int64_t *p10) {
    int limit = 9;
 
    if (!ls) {
        limit = input_digits[i];
    }
 
    int64_t res = 0;
 
    for (int x = 0; x <= limit; ++x) {
        int new_taken = taken || (x == 2) || (x == 3) || (x == 5) || (x == 7);
        int new_i = i + 1;
        int64_t ret;
 
        if (new_i == N) {
            ret = new_taken;
        } else {
            int new_ls = ls || (x != limit);
 
            if (new_ls && new_taken) {
                ret = p10[N - new_i];
            } else {
                int64_t cached = dp[4 * new_i + 2 * new_taken + new_ls];
 
                if (cached != -1) {
                    ret = cached;
                } else {
                    ret = calculate(new_i, new_taken, new_ls, input_digits, N, dp, p10);
                }
            }
        }
 
        res += ret;
    }
 
    dp[4 * i + 2 * taken + ls] = res;
 
    return res;
}

(глобальные переменные мы создаем на стороне питона и передаем как аргументы, также мы преобразовали многомерный массив dp в одномерный, чтобы не париться с указателями)

Дальше нам нужно это скомпилировать, и нужно передать некоторые нестандарные флаги. Обычно исполнение ELF начинается не с main, а с функции _start, которая вызывает функцию __libc_start_main из libc, которая подгружает переменные окружения, argv и многое другое, а только потом запускает main. Нам ничего такого не нужно, поэтому мы это отключаем с помощью флага -nostartfiles, и к тому же без этого компилятор ожидал бы от нас функцию main, которую мы писать не хотим. Также можем добавить -nolibc чтобы не было соблазна вызывать функции оттуда, хотя на результат это не повлияет. И конечно -O3 чтобы получить все оптимизации что есть.

$ gcc inline.c -nostartfiles -nolibc -o inline -O3

Можем проверить что в секции .text у нас лежит одна единственная функция (-d — вывод дизассемблера секции с кодом, -M intel — чтобы получить вывод в намного более приятном интеловском синтаксисе, а не в AT&T):

$ objdump -d -M intel inline
 
inline: file format elf64-x86-64
 
 
Disassembly of section .text:
 
0000000000001000 <calculate>:
    1000:       41 57                   push   r15
    1002:       41 89 f7                mov    r15d,esi
    1005:       41 56                   push   r14
    1007:       41 89 fe                mov    r14d,edi
    100a:       41 55                   push   r13
    100c:       41 89 d5                mov    r13d,edx
    100f:       41 54                   push   r12
    1011:       55                      push   rbp
    1012:       53                      push   rbx
    1013:       48 83 ec 38             sub    rsp,0x38
    1017:       85 d2                   test   edx,edx
    1019:       75 35                   jne    1050 <calculate+0x50>
    101b:       48 63 c7                movsxd rax,edi
    101e:       45 31 e4                xor    r12d,r12d
    1021:       8b 2c 81                mov    ebp,DWORD PTR [rcx+rax*4]
    1024:       85 ed                   test   ebp,ebp
    1026:       79 2d                   jns    1055 <calculate+0x55>
    1028:       43 8d 04 77             lea    eax,[r15+r14*2]
    102c:       41 8d 44 45 00          lea    eax,[r13+rax*2+0x0]
    ...
    ...

Также стоит заметить что наш код не использует никаких глобальных переменных, а значит нам нужна только секция .text из ELF. И еще все инструкции jmp и call используют относительные оффсеты для прыжков, поэтому наш код может быть загружен по любому адресу памяти, что нам и нужно, так как мы не можем полностью контроллировать какой адрес вернет mmap.

Нам осталось сдампить байты нашей функции в отдельный файл (-O binary = вывести сырые байты, -j .text = работать с секцией .text):

$ objcopy -O binary -j .text inline inline_bytes

И записать получившиеся байты внутрь кода в питоне, например так можно получить строку байт чтобы скопировать:

from pprint import pprint
 
with open('inline_bytes', 'rb') as f:
     pprint(f.read())

Какой будет вывод:

(b'AWA\x89\xf7AVA\x89\xfeAUA\x89\xd5ATUSH\x83\xec8\x85\xd2u5Hc\xc7E1'
 b'\xe4\x8b,\x81\x85\xedy-C\x8d\x04wA\x8dDE\x00H\x98M\x89$\xc1H\x83\xc48L'
 b'\x89\xe0[]A\\A]A^A_\xc3\x0f\x1f\x80\x00\x00\x00\x00\xbd\t\x00\x00\x00A\x8d~'
 b'\x01D\x89\xc0H\x8bt$p)\xf8E\x85\xffH\x98A\x0f\x95\xc2L\x8d\x1c\xc6A9\xf8t'
 b'k\x8d\x04?1\xdbE1\xe4\x89D$\x08\x83\xfb\x02t\x19E\x84\xd2u\x14\x83'
 b'\xfb\x07\x0f\x87|\x00\x00\x00\xb8\xa8\x00\x00\x00H\x0f\xa3\xd8sq9\xddu51'
 b'\xd2\xbe\x01\x00\x00\x00E\x85\xedu)\x8bD$\x08\x01\xf0\x8d\x04BH\x98I\x8b'
 b'\x04\xc1H\x83\xf8\xfftdI\x01\xc4\x83\xc3\x019\xdd}\xb3\xe9U\xff\xff\xff\x0f'
 b'\x1fD\x00\x00M\x03#\xeb\xea\x0f\x1f\x001\xc0E1\xe4\xba\xa8\x00'
 b'\x00\x00\x83\xf8\x02t\x10E\x84\xd2u\x0b\x83\xf8\x07w\nH\x0f\xa3\xc2s\x04I'
 b'\x83\xc4\x01\x83\xc0\x019\xc5}\xe0\xe9\x19\xff\xff\xff\x909\xdd\x0f\x95'
 b'\xc2E\x85\xed\x0f\x95\xc01\xf6\t\xc2\x0f\xb6\xd2\xeb\x8b\x0f\x1f@\x00D\x88T$'
 b'/H\x83\xec\x08L\x89\\$(\xfft$xL\x89L$(D\x89D$8H\x89L$ \x89|$\x1c\xe8\xae\xfe'
 b'\xff\xffI\x01\xc4XZ\x8b|$\x0cH\x8bL$\x10D\x8bD$(L\x8bL$\x18L\x8b\\$ D'
 b'\x0f\xb6T$/\xe9M\xff\xff\xff')

И это нужно скопировать, чтобы можно было в codechef послать только один файл:

func_bytes = (b'AWA\x89\xf7AVA\x89\xfeAUA\x89\xd5ATUSH\x83\xec8\x85\xd2u5Hc\xc7E1'
              b'\xe4\x8b,\x81\x85\xedy-C\x8d\x04wA\x8dDE\x00H\x98M\x89$\xc1H\x83\xc48L'
              b'\x89\xe0[]A\\A]A^A_\xc3\x0f\x1f\x80\x00\x00\x00\x00\xbd\t\x00\x00\x00A\x8d~'
              b'\x01D\x89\xc0H\x8bt$p)\xf8E\x85\xffH\x98A\x0f\x95\xc2L\x8d\x1c\xc6A9\xf8t'
              b'k\x8d\x04?1\xdbE1\xe4\x89D$\x08\x83\xfb\x02t\x19E\x84\xd2u\x14\x83'
              b'\xfb\x07\x0f\x87|\x00\x00\x00\xb8\xa8\x00\x00\x00H\x0f\xa3\xd8sq9\xddu51'
              b'\xd2\xbe\x01\x00\x00\x00E\x85\xedu)\x8bD$\x08\x01\xf0\x8d\x04BH\x98I\x8b'
              b'\x04\xc1H\x83\xf8\xfftdI\x01\xc4\x83\xc3\x019\xdd}\xb3\xe9U\xff\xff\xff\x0f'
              b'\x1fD\x00\x00M\x03#\xeb\xea\x0f\x1f\x001\xc0E1\xe4\xba\xa8\x00'
              b'\x00\x00\x83\xf8\x02t\x10E\x84\xd2u\x0b\x83\xf8\x07w\nH\x0f\xa3\xc2s\x04I'
              b'\x83\xc4\x01\x83\xc0\x019\xc5}\xe0\xe9\x19\xff\xff\xff\x909\xdd\x0f\x95'
              b'\xc2E\x85\xed\x0f\x95\xc01\xf6\t\xc2\x0f\xb6\xd2\xeb\x8b\x0f\x1f@\x00D\x88T$'
              b'/H\x83\xec\x08L\x89\\$(\xfft$xL\x89L$(D\x89D$8H\x89L$ \x89|$\x1c\xe8\xae\xfe'
              b'\xff\xffI\x01\xc4XZ\x8b|$\x0cH\x8bL$\x10D\x8bD$(L\x8bL$\x18L\x8b\\$ D'
              b'\x0f\xb6T$/\xe9M\xff\xff\xff')

Вот и все!