├── .gitignore
├── Makefile
├── README.md
├── jitcompiler.py
├── mj.py
├── test-decorator.py
└── test.py


/.gitignore:
--------------------------------------------------------------------------------
1 | *.dSYM
2 | *.o
3 | *.pyc
4 | *.pyo
5 | *.so
6 | GPATH
7 | GRTAGS
8 | GTAGS
9 | 


--------------------------------------------------------------------------------
/Makefile:
--------------------------------------------------------------------------------
1 | test:
2 | 	python mj.py
3 | 	python test.py
4 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
  1 | MiniJIT
  2 | =======
  3 | 
  4 | Contains code for the posts
  5 | 
  6 |   * Writing a basic x86-64 JIT compiler from scratch in stock Python
  7 |     https://csl.name/post/python-jit/
  8 | 
  9 |   * JIT compiling a tiny subset of Python to x86-64 from scratch — in Python
 10 |     https://csl.name/post/python-compiler/
 11 | 
 12 | You need a UNIX/POSIX system and an x86-64 compatible CPU. I've tested this on
 13 | Linux and macOS, using Python 2.7 and 3+
 14 | 
 15 | The ~500 lines of code relies only on standard Python libraries and contains a
 16 | Python bytecode converter, peephole optimizer and x86-64 machine code
 17 | assembler. The code is meant to be simple to understand and pedagogical.
 18 | 
 19 | Finally, there is a decorator that automatically swaps out Python functions
 20 | with native code:
 21 | 
 22 |     >>> from jitcompiler import jit, disassemble
 23 |     >>> @jit
 24 |     ... def foo(a, b): return a + b
 25 |     ... 
 26 |     --- Installing JIT for <function foo at 0x10bc48c08>
 27 |     >>> foo(10, -2)
 28 |     --- JIT-compiling <function foo at 0x10bc48c08>
 29 |     8
 30 |     >>> print(disassemble(foo))
 31 |     0x10bb3c000 48 89 fb       mov rbx, rdi
 32 |     0x10bb3c003 48 89 f0       mov rax, rsi
 33 |     0x10bb3c006 48 01 d8       add rax, rbx
 34 |     0x10bb3c009 c3             ret 
 35 | 
 36 | How to run tests
 37 | ----------------
 38 | 
 39 | The first one patches up some machine code and runs it at runtime
 40 | 
 41 |     $ python mj.py
 42 | 
 43 | The second one JIT compiles Python bytecode to machine code at runtime
 44 | 
 45 |     $ python tests.py
 46 | 
 47 | If you have the `capstone` module installed, it will display an in-memory
 48 | disassembly as well.
 49 | 
 50 | You can also run the decorator test. It defines a function like this
 51 | 
 52 |     import jitcompiler
 53 | 
 54 |     #...
 55 | 
 56 |     @jitcompiler.jit
 57 |     def foo(a, b):
 58 |         return a*a - b*b
 59 | 
 60 | On the first *call* to `foo`, it will be compiled to native code and swap out
 61 | the original Python function. It treats all arguments as signed 64-bit
 62 | integers. If you have the Capstone module installed, it will also print a
 63 | disassembly. To run:
 64 | 
 65 |     $ python test-decorator.py
 66 |     Definition point of foo
 67 | 
 68 |     Installing JIT for <function foo at 0x1f855f0>
 69 | 
 70 |     Calling foo
 71 | 
 72 |     JIT-compiling <function foo at 0x1f855f0>
 73 |     Installed native code for <function foo at 0x1f855f0>
 74 |     Calling function <CFunctionType object at 0x7f867642b600>
 75 |     foo(1, 2) => -3
 76 |     Calling function <CFunctionType object at 0x7f867642b600>
 77 |     foo(2, 3) => -5
 78 | 
 79 |     Disassembly of foo
 80 | 
 81 |     0x7f86765f9000 48 89 fb       mov rbx, rdi
 82 |     0x7f86765f9003 48 89 f8       mov rax, rdi
 83 |     0x7f86765f9006 48 0f af c3    imul rax, rbx
 84 |     0x7f86765f900a 50             push rax
 85 |     0x7f86765f900b 48 89 f3       mov rbx, rsi
 86 |     0x7f86765f900e 48 89 f0       mov rax, rsi
 87 |     0x7f86765f9011 48 0f af c3    imul rax, rbx
 88 |     0x7f86765f9015 48 89 c3       mov rbx, rax
 89 |     0x7f86765f9018 58             pop rax
 90 |     0x7f86765f9019 48 29 d8       sub rax, rbx
 91 |     0x7f86765f901c c3             ret
 92 | 
 93 | If you want to get serious about this
 94 | -------------------------------------
 95 | 
 96 |   * Check out a full-blown assembler library for Python:
 97 |     https://github.com/Maratyszcza/PeachPy
 98 | 
 99 | References
100 | ----------
101 | 
102 |   * Intel assembly manuals:
103 |     https://software.intel.com/en-us/articles/intel-sdm
104 | 
105 |   * x86 Reference: http://ref.x86asm.net/
106 | 
107 |   * Capstone disassembler:
108 |     http://www.capstone-engine.org/lang_python.html
109 | 
110 | License
111 | -------
112 | 
113 | Put in the public domain in 2017 by the author Christian Stigen Larsen
114 | 


--------------------------------------------------------------------------------
/jitcompiler.py:
--------------------------------------------------------------------------------
  1 | """
  2 | JIT compiles a tiny subset of Python to x86-64 machine code.
  3 | Only relies on stock Python.
  4 | 
  5 | Explanation on https://csl.name/post/python-compiler/
  6 | 
  7 | Tested on Python 2.7, 3.4 and 3.6.
  8 | 
  9 | Example Usage
 10 | -------------
 11 | 
 12 |     from jitcompiler import jit
 13 | 
 14 |     @jit
 15 |     def foo(a, b):
 16 |         return a*a - b*b
 17 | 
 18 |     # When foo is called the first time, it will be swapped out with native
 19 |     # code. See the blog posts for details.
 20 | 
 21 | License
 22 | -------
 23 | 
 24 | Written by Christian Stigen Larsen
 25 | Put in the public domain by the author in 2017
 26 | """
 27 | 
 28 | import ctypes
 29 | import dis
 30 | import sys
 31 | 
 32 | # Local include: Provides mmap and related functionality
 33 | import mj
 34 | 
 35 | # Used for compatibility with Python 2.7 and 3+
 36 | PRE36 = sys.version_info[:2] < (3, 6)
 37 | 
 38 | def get_codeobj(function):
 39 |     # NOTE: Seems that __code__ works on Python 2.7 as well now
 40 |     if hasattr(function, "func_code"):
 41 |         return function.func_code
 42 |     else:
 43 |         return function.__code__
 44 | 
 45 | class Assembler(object):
 46 |     """An x86-64 assembler."""
 47 | 
 48 |     def __init__(self, size):
 49 |         self.block = mj.create_block(size)
 50 |         self.index = 0
 51 |         self.size = size
 52 | 
 53 |     @property
 54 |     def raw(self):
 55 |         """Returns machine code as a raw string."""
 56 |         if sys.version_info.major == 2:
 57 |             return "".join(chr(x) for x in self.block[:self.index])
 58 |         else:
 59 |             return bytes(self.block[:self.index])
 60 | 
 61 |     @property
 62 |     def address(self):
 63 |         """Returns address of block in memory."""
 64 |         return ctypes.cast(self.block, ctypes.c_void_p).value
 65 | 
 66 |     def little_endian(self, n):
 67 |         """Converts 64-bit number to little-endian format."""
 68 |         if n is None:
 69 |             n = 0
 70 |         return [(n & (0xff << (i*8))) >> (i*8) for i in range(8)]
 71 | 
 72 |     def registers(self, a, b=None):
 73 |         """Encodes one or two registers for machine code instructions."""
 74 |         order = ("rax", "rcx", "rdx", "rbx", "rsp", "rbp", "rsi", "rdi")
 75 |         enc = order.index(a)
 76 |         if b is not None:
 77 |             enc = enc << 3 | order.index(b)
 78 |         return enc
 79 | 
 80 |     def emit(self, *args):
 81 |         """Writes machine code to memory block."""
 82 |         for code in args:
 83 |             self.block[self.index] = code
 84 |             self.index += 1
 85 | 
 86 |     def ret(self, a, b):
 87 |         self.emit(0xc3)
 88 | 
 89 |     def push(self, a, _):
 90 |         self.emit(0x50 | self.registers(a))
 91 | 
 92 |     def pop(self, a, _):
 93 |         self.emit(0x58 | self.registers(a))
 94 | 
 95 |     def imul(self, a, b):
 96 |         self.emit(0x48, 0x0f, 0xaf, 0xc0 | self.registers(a, b))
 97 | 
 98 |     def add(self, a, b):
 99 |         self.emit(0x48, 0x01, 0xc0 | self.registers(b, a))
100 | 
101 |     def sub(self, a, b):
102 |         self.emit(0x48, 0x29, 0xc0 | self.registers(b, a))
103 | 
104 |     def neg(self, a, _):
105 |         self.emit(0x48, 0xf7, 0xd8 | self.registers(a))
106 | 
107 |     def mov(self, a, b):
108 |         self.emit(0x48, 0x89, 0xc0 | self.registers(b, a))
109 | 
110 |     def immediate(self, a, number):
111 |         self.emit(0x48, 0xb8 | self.registers(a), *self.little_endian(number))
112 | 
113 | 
114 | class Compiler(object):
115 |     """Compiles Python bytecode to intermediate representation (IR)."""
116 |     def __init__(self, bytecode, constants):
117 |         self.bytecode = bytecode
118 |         self.constants = constants
119 |         self.index = 0
120 | 
121 |     def fetch(self):
122 |         byte = self.bytecode[self.index]
123 |         self.index += 1
124 |         return byte
125 | 
126 |     def decode(self):
127 |         opcode = self.fetch()
128 |         opname = dis.opname[opcode]
129 | 
130 |         if opname.startswith(("UNARY", "BINARY", "INPLACE", "RETURN")):
131 |             argument = None
132 |             if not PRE36:
133 |                 self.fetch()
134 |         else:
135 |             argument = self.fetch()
136 |             if PRE36:
137 |                 argument |= self.fetch() << 8
138 | 
139 |         return opname, argument
140 | 
141 |     def variable(self, number):
142 |         # AMD64 argument passing order for our purposes.
143 |         order = ("rdi", "rsi", "rdx", "rcx")
144 |         return order[number]
145 | 
146 |     def compile(self):
147 |         while self.index < len(self.bytecode):
148 |             op, arg = self.decode()
149 | 
150 |             if op == "LOAD_FAST":
151 |                 yield "push", self.variable(arg), None
152 | 
153 |             elif op == "STORE_FAST":
154 |                 yield "pop", "rax", None
155 |                 yield "mov", self.variable(arg), "rax"
156 | 
157 |             elif op == "LOAD_CONST":
158 |                 value = self.constants[arg]
159 |                 if value is None:
160 |                     value = 0
161 |                 yield "immediate", "rax", value
162 |                 yield "push", "rax", None
163 | 
164 |             elif op == "BINARY_MULTIPLY":
165 |                 yield "pop", "rax", None
166 |                 yield "pop", "rbx", None
167 |                 yield "imul", "rax", "rbx"
168 |                 yield "push", "rax", None
169 | 
170 |             elif op in ("BINARY_ADD", "INPLACE_ADD"):
171 |                 yield "pop", "rax", None
172 |                 yield "pop", "rbx", None
173 |                 yield "add", "rax", "rbx"
174 |                 yield "push", "rax", None
175 | 
176 |             elif op in ("BINARY_SUBTRACT", "INPLACE_SUBTRACT"):
177 |                 yield "pop", "rbx", None
178 |                 yield "pop", "rax", None
179 |                 yield "sub", "rax", "rbx"
180 |                 yield "push", "rax", None
181 | 
182 |             elif op == "UNARY_NEGATIVE":
183 |                 yield "pop", "rax", None
184 |                 yield "neg", "rax", None
185 |                 yield "push", "rax", None
186 | 
187 |             elif op == "RETURN_VALUE":
188 |                 yield "pop", "rax", None
189 |                 yield "ret", None, None
190 |             else:
191 |                 raise NotImplementedError(op)
192 | 
193 | def optimize(ir):
194 |     """Performs peephole optimizations on the IR."""
195 |     def fetch(n):
196 |         if n < len(ir):
197 |             return ir[n]
198 |         else:
199 |             return None, None, None
200 | 
201 |     index = 0
202 |     while index < len(ir):
203 |         op1, a1, b1 = fetch(index)
204 |         op2, a2, b2 = fetch(index + 1)
205 |         op3, a3, b3 = fetch(index + 2)
206 |         op4, a4, b4 = fetch(index + 3)
207 | 
208 |         # Remove nonsensical moves
209 |         if op1 == "mov" and a1 == b1:
210 |             index += 1
211 |             continue
212 | 
213 |         # Translate
214 |         #    mov rsi, rax
215 |         #    mov rbx, rsi
216 |         # to mov rbx, rax
217 |         if op1 == op2 == "mov" and a1 == b2:
218 |             index += 2
219 |             yield "mov", a2, b1
220 |             continue
221 | 
222 |         # Short-circuit push x/pop y
223 |         if op1 == "push" and op2 == "pop":
224 |             index += 2
225 |             yield "mov", a2, a1
226 |             continue
227 | 
228 |         # Same as above, but with an in-between instruction
229 |         if op1 == "push" and op3 == "pop" and op2 not in ("push", "pop"):
230 |             # Only do this if a3 is not modified in the middle instruction. An
231 |             # obvious improvement would be to allow an arbitrary number of
232 |             # in-between instructions.
233 |             if a2 != a3:
234 |                 index += 3
235 |                 yield "mov", a3, a1
236 |                 yield op2, a2, b2
237 |                 continue
238 | 
239 |         # Same as above, but with one in-between instruction.
240 |         # TODO: Generalize this, then remove the previous two
241 |         if (op1 == "push" and op4 == "pop" and op2 not in ("push", "pop") and
242 |                 op3 not in ("push", "pop")):
243 |             if a2 != a4 and a3 != a4:
244 |                 index += 4
245 |                 yield "mov", a4, a1
246 |                 yield op2, a2, b2
247 |                 yield op3, a3, b3
248 |                 continue
249 | 
250 |         index += 1
251 |         yield op1, a1, b1
252 | 
253 | def print_ir(ir):
254 |     for instruction in ir:
255 |         op, args = instruction[0], instruction[1:]
256 |         args = filter(lambda x: x is not None, args)
257 |         print("  %-6s %s" % (op, ", ".join(map(str, args))))
258 | 
259 | def compile_native(function, verbose=True):
260 |     """Compiles a branchless Python function to native x86-64 machine code.
261 | 
262 |     Returns:
263 |         A tuple consisting of a callable Python function bound to the native
264 |         code, and the assembler used to create it (mostly for disassembly
265 |         purposes).
266 |     """
267 |     if verbose:
268 |         print("Python disassembly:")
269 |         dis.dis(function)
270 |         print("")
271 | 
272 |     codeobj = get_codeobj(function)
273 |     if verbose:
274 |         print("Bytecode: %r" % codeobj.co_code)
275 |         print("")
276 | 
277 |     if verbose:
278 |         print("Intermediate code:")
279 |     constants = codeobj.co_consts
280 | 
281 |     python_bytecode = list(codeobj.co_code)
282 | 
283 |     if sys.version_info.major == 2:
284 |         python_bytecode = map(ord, codeobj.co_code)
285 | 
286 |     ir = Compiler(python_bytecode, constants).compile()
287 |     ir = list(ir)
288 |     if verbose:
289 |         print_ir(ir)
290 |         print("")
291 | 
292 |     if verbose:
293 |         print("Optimization:")
294 |     while True:
295 |         optimized = list(optimize(ir))
296 |         reduction = len(ir) - len(optimized)
297 |         ir = optimized
298 |         if verbose:
299 |             print("  - removed %d instructions" % reduction)
300 |         if not reduction:
301 |             break
302 |     if verbose:
303 |         print_ir(ir)
304 |         print("")
305 | 
306 |     # Compile to native code
307 |     assembler = Assembler(mj.PAGESIZE)
308 |     for name, a, b in ir:
309 |         emit = getattr(assembler, name)
310 |         emit(a, b)
311 | 
312 |     # Make block executable and read-only
313 |     mj.make_executable(assembler.block, assembler.size)
314 | 
315 |     argcount = codeobj.co_argcount
316 | 
317 |     if argcount == 0:
318 |         signature = ctypes.CFUNCTYPE(None)
319 |     else:
320 |         # Assume all arguments are 64-bit
321 |         signature = ctypes.CFUNCTYPE(*[ctypes.c_int64] * argcount)
322 | 
323 |     signature.restype = ctypes.c_int64
324 |     return signature(assembler.address), assembler
325 | 
326 | def jit(function):
327 |     """Decorator that JIT-compiles function to native code on first call.
328 | 
329 |     Use this on non-class functions, because our compiler does not support
330 |     objects (rather, it does not support the attr bytecode instructions).
331 | 
332 |     Also, only works on branchless Python functions that only perform
333 |     arithmetic on signed, 64-bit integers.
334 | 
335 |     Example:
336 |         @jit
337 |         def foo(a, b):
338 |             return a*a - b*b
339 |     """
340 |     print("--- Installing JIT for %s" % function)
341 | 
342 |     def frontend(*args, **kw):
343 |         if not hasattr(frontend, "function"):
344 |             try:
345 |                 print("--- JIT-compiling %s" % function)
346 |                 native, asm = compile_native(function, verbose=False)
347 |                 native.raw = asm.raw
348 |                 native.address = asm.address
349 |                 frontend.function = native
350 |             except Exception as e:
351 |                 frontend.function = function # fallback to Python
352 |                 print("--- Could not compile %s: %s: %s" % (function.__name__,
353 |                     e.__class__.__name__, e))
354 |         return frontend.function(*args, **kw)
355 |     return frontend
356 | 
357 | def disassemble(function):
358 |     """Returns disassembly string of natively compiled function.
359 | 
360 |     Requires the Capstone module."""
361 |     if hasattr(function, "function"):
362 |         function = function.function
363 | 
364 |     def hexbytes(raw):
365 |         return "".join("%02x " % b for b in raw)
366 | 
367 |     try:
368 |         import capstone
369 | 
370 |         out = ""
371 |         md = capstone.Cs(capstone.CS_ARCH_X86, capstone.CS_MODE_64)
372 | 
373 |         for i in md.disasm(function.raw, function.address):
374 |             out += "0x%x %-15s%s %s\n" % (i.address, hexbytes(i.bytes), i.mnemonic, i.op_str)
375 |             if i.mnemonic == "ret":
376 |                 break
377 | 
378 |         return out
379 |     except ImportError:
380 |         print("You need to install the Capstone module for disassembly")
381 |         raise
382 | 


--------------------------------------------------------------------------------
/mj.py:
--------------------------------------------------------------------------------
  1 | """
  2 | Provides a way to generate machine code and bind it to callable Python
  3 | functions at runtime.
  4 | 
  5 | You need a UNIX system with mmap, mprotect and so on. Tested on macOS and
  6 | Linux.
  7 | 
  8 | See https://csl.name/post/python-jit/ for a write-up on how everything works!
  9 | 
 10 | Written by Christian Stigen Larsen
 11 | """
 12 | 
 13 | import ctypes
 14 | import ctypes.util
 15 | import mmap as MMAP
 16 | import os
 17 | import sys
 18 | 
 19 | # Load the C standard library
 20 | libc = ctypes.CDLL(ctypes.util.find_library("c"))
 21 | 
 22 | # A few constants
 23 | MAP_FAILED = -1 # voidptr actually
 24 | 
 25 | # Set up strerror
 26 | strerror = libc.strerror
 27 | strerror.argtypes = [ctypes.c_int]
 28 | strerror.restype = ctypes.c_char_p
 29 | 
 30 | # Get pagesize
 31 | PAGESIZE = os.sysconf(os.sysconf_names["SC_PAGESIZE"])
 32 | 
 33 | # 8-bit unsigned pointer type
 34 | c_uint8_p = ctypes.POINTER(ctypes.c_uint8)
 35 | 
 36 | # Setup mmap
 37 | mmap = libc.mmap
 38 | mmap.argtypes = [ctypes.c_void_p,
 39 |                  ctypes.c_size_t,
 40 |                  ctypes.c_int,
 41 |                  ctypes.c_int,
 42 |                  ctypes.c_int,
 43 |                  # Below is actually off_t, which is 64-bit on macOS
 44 |                  ctypes.c_int64]
 45 | mmap.restype = c_uint8_p
 46 | 
 47 | # Setup munmap
 48 | munmap = libc.munmap
 49 | munmap.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
 50 | munmap.restype = ctypes.c_int
 51 | 
 52 | # Set mprotect
 53 | mprotect = libc.mprotect
 54 | mprotect.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]
 55 | mprotect.restype = ctypes.c_int
 56 | 
 57 | def create_block(size):
 58 |     """Allocated a block of memory using mmap."""
 59 |     ptr = mmap(0, size, MMAP.PROT_WRITE | MMAP.PROT_READ,
 60 |             MMAP.MAP_PRIVATE | MMAP.MAP_ANONYMOUS, 0, 0)
 61 | 
 62 |     if ptr == MAP_FAILED:
 63 |         raise RuntimeError(strerror(ctypes.get_errno()))
 64 | 
 65 |     return ptr
 66 | 
 67 | def make_executable(block, size):
 68 |     """Marks mmap'ed memory block as read-only and executable."""
 69 |     if mprotect(block, size, MMAP.PROT_READ | MMAP.PROT_EXEC) != 0:
 70 |         raise RuntimeError(strerror(ctypes.get_errno()))
 71 | 
 72 | def destroy_block(block, size):
 73 |     """Deallocated previously mmapped block."""
 74 |     if munmap(block, size) == -1:
 75 |         raise RuntimeError(strerror(ctypes.get_errno()))
 76 |     del block
 77 | 
 78 | def make_multiplier(block, multiplier):
 79 |     """JIT-compiles a function that multiplies its RDX argument with an
 80 |     unsigned 64-bit constant."""
 81 |     if multiplier > (2**64-1) or multiplier < 0:
 82 |         raise ValueError("Multiplier does not fit in unsigned 64-bit integer")
 83 | 
 84 |     # This function encodes the disassembly of multiply.c, which you can see
 85 |     # with the command `make dis`. It may be different on your CPU, so adjust
 86 |     # to match.
 87 |     #
 88 |     #   48 b8 ed ef be ad de    movabs $0xdeadbeefed,%rax
 89 |     #   00 00 00
 90 |     #   48 0f af c7             imul   %rdi,%rax
 91 |     #   c3                      retq
 92 | 
 93 |     # Encoding of: movabs <multiplier>, rax
 94 |     block[0] = 0x48
 95 |     block[1] = 0xb8
 96 | 
 97 |     # Little-endian encoding of multiplier
 98 |     block[2] = (multiplier & 0x00000000000000ff) >>  0
 99 |     block[3] = (multiplier & 0x000000000000ff00) >>  8
100 |     block[4] = (multiplier & 0x0000000000ff0000) >> 16
101 |     block[5] = (multiplier & 0x00000000ff000000) >> 24
102 |     block[6] = (multiplier & 0x000000ff00000000) >> 32
103 |     block[7] = (multiplier & 0x0000ff0000000000) >> 40
104 |     block[8] = (multiplier & 0x00ff000000000000) >> 48
105 |     block[9] = (multiplier & 0xff00000000000000) >> 56
106 | 
107 |     # Encoding of: imul rdi, rax
108 |     block[10] = 0x48
109 |     block[11] = 0x0f
110 |     block[12] = 0xaf
111 |     block[13] = 0xc7
112 | 
113 | 
114 |     # Encoding of: retq
115 |     block[14] = 0xc3
116 | 
117 |     # Return a ctypes function with the right prototype
118 |     function = ctypes.CFUNCTYPE(ctypes.c_uint64)
119 |     function.restype = ctypes.c_uint64
120 |     return function
121 | 
122 | def main():
123 |     # Fetch the constant to multiply with on the command line. If not
124 |     # specified, use the default value of 11.
125 |     if len(sys.argv) > 1:
126 |         arg = int(sys.argv[1])
127 |     else:
128 |         arg = 11
129 | 
130 |     print("Pagesize: %d" % PAGESIZE)
131 | 
132 |     print("Allocating one page of memory")
133 |     block = create_block(PAGESIZE)
134 | 
135 |     print("JIT-compiling a native mul-function w/arg %d" % arg)
136 |     function_type = make_multiplier(block, arg)
137 | 
138 |     print("Making function block executable")
139 |     make_executable(block, PAGESIZE)
140 |     mul = function_type(ctypes.cast(block, ctypes.c_void_p).value)
141 | 
142 |     print("Testing function")
143 |     for i in range(10):
144 |         expected = i*arg
145 |         actual = mul(i)
146 |         print("%-4s mul(%d) = %d" % ("OK" if actual == expected else "FAIL", i,
147 |             actual))
148 | 
149 |     print("Deallocating function")
150 |     destroy_block(block, PAGESIZE)
151 | 
152 |     # Unbind local variables
153 |     del block
154 |     del mul
155 | 
156 | if __name__ == "__main__":
157 |     main()
158 | 


--------------------------------------------------------------------------------
/test-decorator.py:
--------------------------------------------------------------------------------
 1 | """
 2 | Shows how to use the jitcompiler.jit decorator to automatically compile the
 3 | function to native code on the first call.
 4 | """
 5 | 
 6 | import jitcompiler
 7 | 
 8 | print("Definition point of foo\n")
 9 | 
10 | @jitcompiler.jit
11 | def foo(a, b):
12 |     return a*a - b*b
13 | 
14 | def test(a, b):
15 |     result = foo(a, b)
16 |     print("foo(%d, %d) => %d" % (a, b, result))
17 |     assert(result == (a*a - b*b))
18 | 
19 | print("\nCalling foo\n")
20 | test(1, 2)
21 | test(2, 3)
22 | 
23 | print("\nDisassembly of foo\n")
24 | print(jitcompiler.disassemble(foo))
25 | 


--------------------------------------------------------------------------------
/test.py:
--------------------------------------------------------------------------------
  1 | """
  2 | Tests for: JIT compiles a tiny subset of Python to x86-64 machine code.
  3 | 
  4 | Explanation on https://csl.name/post/python-compiler/
  5 | 
  6 | Written by Christian Stigen Larsen
  7 | Put in the public domain by the author, 2017
  8 | """
  9 | 
 10 | import random
 11 | import jitcompiler
 12 | import sys
 13 | 
 14 | try:
 15 |     import capstone
 16 | except ImportError:
 17 |     pass
 18 | 
 19 | good = True
 20 | 
 21 | def test_function(original, compiled, tests=10, minval=-999, maxval=999):
 22 |     for n in range(tests):
 23 |         # Create random arguments
 24 |         argcount = jitcompiler.get_codeobj(original).co_argcount
 25 |         args = [random.randint(minval, maxval) for x in range(argcount)]
 26 | 
 27 |         # Run original and compiled functions
 28 |         expected = original(*args)
 29 |         actual = compiled(*args)
 30 |         ok = (expected == actual)
 31 | 
 32 |         global good
 33 |         if not ok:
 34 |             good = False
 35 | 
 36 |         print("  %-4s %-16s => %10d, expected %10d" % (
 37 |             "OK" if ok else "FAIL",
 38 |             "(%s)" % ", ".join("%4d" % d for d in args),
 39 |             actual,
 40 |             expected))
 41 | 
 42 | def test(function):
 43 |     print("")
 44 |     print("=== Function %s ===" % function.__name__)
 45 |     print("")
 46 | 
 47 |     native, asm = jitcompiler.compile_native(function)
 48 | 
 49 |     try:
 50 |         print("Native code:")
 51 |         md = capstone.Cs(capstone.CS_ARCH_X86, capstone.CS_MODE_64)
 52 |         for i in md.disasm(asm.raw, asm.address):
 53 |             print("  0x%x:\t%s\t%s" % (i.address, i.mnemonic, i.op_str))
 54 |             if i.mnemonic == "ret":
 55 |                 break
 56 |         print("")
 57 |     except NameError:
 58 |         pass
 59 | 
 60 |     test_function(function, native)
 61 | 
 62 | if __name__ == "__main__":
 63 |     def example0(n):
 64 |         return n
 65 | 
 66 |     def example1(n):
 67 |         return n*101
 68 | 
 69 |     def example2(a, b):
 70 |         return a*a + b*b
 71 | 
 72 |     def example3(a):
 73 |         b = a*101
 74 |         return b + a + 2
 75 | 
 76 |     def example4(a, b, c):
 77 |         return a*a + 2*a*b + c
 78 | 
 79 |     def example5(n):
 80 |         n -= 10
 81 |         return n
 82 | 
 83 |     def example6(a, b):
 84 |         return a*a - b*b
 85 | 
 86 |     def example7(a, b, c):
 87 |         return (a+c)*b - a*a*(a-c-b)-b*2+(c*(2+3*a*b-c*a)-3*c)
 88 | 
 89 |     def foo(a, b):
 90 |         return a*a - b*b
 91 | 
 92 |     def bar(n):
 93 |         return n+123456789
 94 | 
 95 |     test(example0)
 96 |     test(example1)
 97 |     test(example2)
 98 |     test(example3)
 99 |     test(example4)
100 |     test(example5)
101 |     test(example6)
102 |     #test(example7) # works, but produces too much output
103 |     test(foo)
104 |     test(bar)
105 | 
106 |     if not good:
107 |         print("")
108 |         print("One or more errors occurred.")
109 |         sys.exit(1)
110 | 


--------------------------------------------------------------------------------