pyasm 0.2 - dynamic x86 assembler for python

[Michael Spencer]

JanC wrote:
> [an example of using decorators to control pyasm]

Another (perhaps wacky) approach would be to change the assembler source syntax 
enough to make it legal Python - in particular, this means parenthesizing the 
arguments - then it can just be stored in-line with other Python source.  This 
has the additional benefit that one could write support functions to enable the 
source to be executed interactively in Python.

The following example uses the CPython opcodes, represented as Python functions. 
  Python control structures 'while' and 'if' are used as assembler directives 
for flow.

Michael

  >>> import ackermann
  >>> Ackermann = assemble(ackermann.ackSrc)
  [snip assembler output]
  >>> Ackermann
  <function ackSrc at 0x0185A570>
  >>> Ackermann(3,8)
  2045


# ackermann.py  --------------------------------------------------

def ackSrc(m,n):
     "Compute Ackermann's function on a stack"
     # Can be assembled to Python bytecode, or (not implemented yet)
     # executed in Python with suitable support functions


     LOAD_CONST("Return")
     LOAD_FAST(m)
     LOAD_FAST(n)

     while condition(ROT_TWO(), DUP_TOP(), LOAD_CONST("Return"), COMPARE_OP("!=")):

         if condition(POP_TOP(), DUP_TOP(), LOAD_CONST(0), COMPARE_OP("==")):
             POP_TOP()
             POP_TOP()
             LOAD_CONST(1)
             BINARY_ADD()

         else:
             if condition(POP_TOP(), ROT_TWO(), DUP_TOP(), LOAD_CONST(0), 
COMPARE_OP("==")):
                 POP_TOP()
                 POP_TOP()
                 LOAD_CONST(1)
                 BINARY_SUBTRACT()
                 LOAD_CONST(1)
             else:
                 POP_TOP()
                 DUP_TOP()
                 LOAD_CONST(1)
                 BINARY_SUBTRACT()
                 ROT_THREE()
                 POP_TOP()
                 DUP_TOP()
                 LOAD_CONST(1)
                 BINARY_SUBTRACT()
                 ROT_THREE()
                 ROT_TWO()
     POP_TOP()
     POP_TOP()
     return

# ByteCode.py --------------------------------------------------

"""Python ByteCode Assembler

Author: Michael Spencer
Version: 0 - Experiment
3/11/2005

Example usage:
  >>> import ackermann
  >>> Ackermann = assemble(ackermann.ackSrc)
  [snip assembler output]
  >>> Ackermann
  <function ackSrc at 0x0185A570>
  >>> Ackermann(3,8)
  2045
"""


import dis
import compiler
import compiler.ast as ast
opmap = dis.opmap
import new
import inspect


class AbstractVisitor(object):
     """Standard depth-first AST walker - dispatches to methods
         based on Node class name"""
     def __init__(self):
         self._cache = {} # dispatch table

     def visit(self, node,**kw):
         #print "Visiting: %s" % node.__class__

         if node is None: return None
         cls = node.__class__
         meth = self._cache.setdefault(cls,
             getattr(self,'visit'+cls.__name__,self.default))
         return meth(node, **kw)

     def default(self, node, **kw):
         for child in node.getChildNodes():
             self.visit(child, **kw)
     visitExpression = default


class Assembler(AbstractVisitor):
     """Python bytecode assembler"""

     def __init__(self):
         self._cache = {} # dispatch table
         self.i = 0       # Bytecode instruction counter

         self.co_varnames = []
         self.co_consts = []
         self.jumptable = {}
         self.co_codelst = []

     def emit(self, funcname, arg = None):
         i = self.i
         try:
             opcode = opmap[funcname]
         except KeyError:
             raise SyntaxError, "Unknown operation: %s" % funcname
         self.co_codelst.append(opcode)
         if opcode > dis.HAVE_ARGUMENT:
             print "%4s %4s  %s %s" % (i, opcode, funcname.ljust(20), arg)
             self.co_codelst.extend(self._getlohi(arg))
             self.i = i + 3
         else:
             print "%4s %4s  %s" % (i, opcode, funcname.ljust(20))

             self.i = i + 1

     def getcodestring(self):
         self._resolvejumps()
         return "".join(map(chr, self.co_codelst))

     def getcode(self):
         return new.code(self.co_argcount, # argcount
                     self.co_argcount,     # nlocals
                     10000,                # stacksize
                     67,                   # flags
                     self.getcodestring(), # codestring
                     tuple(self.co_consts),  # constants
                     tuple(self.co_varnames),    # names
                     tuple(self.co_varnames),    # varnames
                     "assembly",                 # filename
                     self.co_name,               # name
                     0,                          # firstlineno
                     ""                          # lnotab
                     )

     def _getlohi(self, arg):
         if isinstance(arg, int):
             return arg % 256, arg / 256
         else:
             return None,None

     def _resolvejumps(self):
         for origin, dest in self.jumptable.iteritems():
             self.co_codelst[origin+1:origin+3] = self._getlohi(dest - origin - 3)

     def visitFunction(self, node, **kw):

         self.co_name = node.name
         self.co_argcount = len(node.argnames)
         self.co_varnames.extend(node.argnames)

         print "def %s(%s)" % (node.name, node.argnames)
         self.visit(node.code)



     def visitCallFunc(self,node,**kw):
         funcname = node.node.name
         if funcname == 'COMPARE_OP': # Special case
             try:
                 comptype = node.args[0].value
                 arg = self.comparisons.index(comptype)
             except ValueError:
                 raise SyntaxError, "Unknown comparison %s" % comptype
         else:
             args = [self.visit(arg) for arg in node.args]
             if args:
                 arg = args[0]
             else:
                 arg = None

         if funcname != "condition": # special case, emits no code
             self.emit(funcname, arg)

     comparisons = list(dis.cmp_op)

     def visitConst(self, node, **kw):
         val = node.value
         try:
             return self.co_consts.index(val)
         except ValueError:
             self.co_consts.append(val)
             return len(self.co_consts)-1

     def visitName(self, node, **kw):
         name = node.name
         try:
             return self.co_varnames.index(name)
         except ValueError:
             self.co_varnames.append(name)
             return len(self.co_consts)-1

     def visitIf(self, node, **kw):
         "node.tests = [(condition, stmt),...] node.else_ = stmt"

         self.visit(node.tests[0][0])  # Get the predicate suite

         jumpbase = self.i
         self.emit("JUMP_IF_FALSE")
         self.visit(node.tests[0][1])

         if node.else_:
             elsebase = self.i
             self.emit("JUMP_FORWARD")
             self.jumptable[jumpbase] = self.i
             print ">> comefrom %s" % jumpbase

             self.visit(node.else_)
             self.jumptable[elsebase] = self.i

             print ">> comefrom %s" % elsebase
         else:
             self.jumptable[jumpbase] = self.i
             print ">> comefrom %s" % jumpbase

     def visitReturn(self, node, **kw):
         self.emit("RETURN_VALUE")


     def visitWhile(self, node, **kw):

         loopstart = self.i
         self.visit(node.test)
         looptest = self.i
         self.emit("JUMP_IF_FALSE")

         self.visit(node.body)
         self.emit("JUMP_ABSOLUTE",loopstart)
         print ">> comefrom %s" % looptest
         self.jumptable[looptest] = self.i


def assemble(source):
     if type(source) == type(assemble):
             source = inspect.getsource(source)

     tree = compiler.parse(source)
     bc = Assembler()
     bc.visit(tree)
     return new.function(bc.getcode(),globals())