root/tags/1.4.0/codewalk.py

"""Bytecode visitors, including rewriters and decompilers.

This work, including the source code, documentation
and related data, is placed into the public domain.

The orginal author is Robert Brewer, Amor Ministries.

THIS SOFTWARE IS PROVIDED AS-IS, WITHOUT WARRANTY
OF ANY KIND, NOT EVEN THE IMPLIED WARRANTY OF
MERCHANTABILITY. THE AUTHOR OF THIS SOFTWARE
ASSUMES _NO_ RESPONSIBILITY FOR ANY CONSEQUENCE
RESULTING FROM THE USE, MODIFICATION, OR
REDISTRIBUTION OF THIS SOFTWARE.

"""

from opcode import cmp_op, opname, opmap, HAVE_ARGUMENT
import operator
import sets
from types import CodeType, FunctionType


def named_opcodes(bits):
    """Change initial numeric opcode bits to their named equivalents."""
    bitnums = []
    bits = iter(bits)
    for x in bits:
        bitnums.append(opname[x])
        if x >= HAVE_ARGUMENT:
            try:
                bitnums.append(bits.next())
                bitnums.append(bits.next())
            except StopIteration:
                break
    return bitnums

def numeric_opcodes(bits):
    """Change named opcode bits to their numeric equivalents."""
    bitnums = []
    for x in bits:
        if isinstance(x, basestring):
            x = opmap[x]
        bitnums.append(x)
    return bitnums

_deref_bytecode = numeric_opcodes(['LOAD_DEREF', 0, 0, 'RETURN_VALUE'])
# CodeType(argcount, nlocals, stacksize, flags, codestring, constants,
#          names, varnames, filename, name, firstlineno,
#          lnotab[, freevars[, cellvars]])
_derefblock = CodeType(0, 0, 1, 3, ''.join(map(chr, _deref_bytecode)),
                       (None,), ('cell',), (), '', '', 2, '', ('cell',))
def deref_cell(cell):
    # FunctionType(code, globals[, name[, argdefs[, closure]]])
    return FunctionType(_derefblock, {}, "", (), (cell,))()

def make_closure(*args):
    def inner():
        args
    return inner.func_closure


binary_operators = {'BINARY_POWER': operator.pow,
                    'BINARY_MULTIPLY': operator.mul,
                    'BINARY_DIVIDE': operator.div,
                    'BINARY_FLOOR_DIVIDE': operator.floordiv,
                    'BINARY_TRUE_DIVIDE': operator.truediv,
                    'BINARY_MODULO': operator.mod,
                    'BINARY_ADD': operator.add,
                    'BINARY_SUBTRACT': operator.sub,
                    'BINARY_SUBSCR': operator.getitem,
                    'BINARY_LSHIFT': operator.lshift,
                    'BINARY_RSHIFT': operator.rshift,
                    'BINARY_AND': operator.and_,
                    'BINARY_XOR': operator.xor,
                    'BINARY_OR': operator.or_,
                    }
inplace_operators = dict([('INPLACE_' + k.split('_')[1], v)
                          for k, v in binary_operators.iteritems()
                          if k not in ('BINARY_SUBSCR',)
                          ])

binary_repr = {'BINARY_POWER': '**',
               'BINARY_MULTIPLY': '*',
               'BINARY_DIVIDE': '/',
               'BINARY_FLOOR_DIVIDE': '//',
               'BINARY_TRUE_DIVIDE': '/',
               'BINARY_MODULO': '%',
               'BINARY_ADD': '+',
               'BINARY_SUBTRACT': '-',
               'BINARY_LSHIFT': '<<',
               'BINARY_RSHIFT': '>>',
               'BINARY_AND': '&',
               'BINARY_XOR': '^',
               'BINARY_OR': '|',
               }

inplace_repr = dict([('INPLACE_' + k.split('_')[1], v + '=')
                     for k, v in binary_repr.iteritems()])

comparisons = {'<': operator.lt,
               '<=': operator.le,
               '==': operator.eq,
               '!=': operator.ne,
               '>': operator.gt,
               '>=': operator.gt,
               'in': operator.contains,
               'not in': lambda x, y: not x in y,
               'is': operator.is_,
               'is not': operator.is_not,
               }

_co_code_attrs = ['co_argcount', 'co_cellvars', 'co_code', 'co_consts',
                  'co_filename', 'co_firstlineno', 'co_flags', 'co_freevars',
                  'co_lnotab', 'co_name', 'co_names', 'co_nlocals',
                  'co_stacksize', 'co_varnames']


class Visitor(object):
    """Visitor(obj=function, code object, string, or list of opcodes)."""
    
    def __init__(self, obj):
        self.verbose = False
        
        # Distill supplied 'obj' arg to a code block string.
        if isinstance(obj, FunctionType):
            obj = obj.func_code
        try:
            obj = obj.co_code
        except AttributeError:
            pass
        
        # Map the code block string to a list of opcode numbers.
        if isinstance(obj, basestring):
            obj = map(ord, obj)
        elif isinstance(obj, list):
            obj = obj[:]
        
        self._bytecode = obj
    
    def debug(self, *messages):
        if self.verbose:
            for term in messages:
                print term,
        else:
            pass
    
    def walk(self):
        self.cursor = 0
        b = self._bytecode
        self.debug("\nWALKING: ", b)
        while self.cursor < len(b):
            self.debug("\n", self.cursor)
            
            op = b[self.cursor]
            self.cursor += 1
            if op >= HAVE_ARGUMENT:
                lo = b[self.cursor]
                self.cursor += 1
                hi = b[self.cursor]
                self.cursor += 1
                args = (lo, hi)
            else:
                args = ()
            
            self.debug("visit (%s, %s)" % (op, repr(args)))
            self.visit_instruction(op, *args)
            
            name = 'visit_' + opname[op]
            name = name.replace('+', '_PLUS_')
            handler = getattr(self, name, None)
            if handler:
                self.debug("=> %s%s" % (name[6:], repr(args)))
                handler(*args)
    
    def visit_instruction(self, op, lo=None, hi=None):
        pass


class JumpCodeAdjuster(Visitor):
    """JumpCodeAdjuster(obj=[func|co|str|list], start, end, newlength).
    
    Adjusts jump codes if their target is affected by bytecode changes.
    
    start, end: The range of the original bytecode in question.
    newlength: Length of the codes which overwrote bytecode[start:end].
    """
    
    def __init__(self, obj, start, end, newlength):
        Visitor.__init__(self, obj)
        self.start = start
        self.end = end
        self.offset = newlength - (end - start)
    
    def bytecode(self):
        """Walk self and return new bytecode."""
        self.walk()
        return self.newcode
    
    def walk(self):
        if self.offset == 0:
            # Avoid costly walk if no changes will be made.
            self.newcode = self._bytecode
        else:
            self.newcode = []
            Visitor.walk(self)
    
    def visit_instruction(self, op, lo=None, hi=None):
        append = self.newcode.append
        append(op)
        if lo is not None:
            append(lo)
        if hi is not None:
            append(hi)
    
    def visit_CONTINUE_LOOP(self, lo, hi):
        self.visit_JUMP_ABSOLUTE(lo, hi)
    
    def visit_JUMP_ABSOLUTE(self, lo, hi):
        target = lo + (hi << 8)
        if target > self.start:
            pos = target + self.offset
            self.newcode[-2:] = [pos & 0xFF, pos >> 8]
    
    def visit_JUMP_FORWARD(self, lo, hi):
        delta = lo + (hi << 8)
        target = self.cursor + delta
        if self.cursor < self.end and target > self.start:
            pos = (target + self.offset) - self.cursor
            self.newcode[-2:] = [pos & 0xFF, pos >> 8]
    
    def visit_JUMP_IF_FALSE(self, lo, hi):
        self.visit_JUMP_FORWARD(lo, hi)
    
    def visit_JUMP_IF_TRUE(self, lo, hi):
        self.visit_JUMP_FORWARD(lo, hi)


def safe_tuple(seq):
    """Force func_code attributes to tuples of strings.
    
    Many of the func_code attributes must take tuples, not lists,
    and *cannot* accept unicode items--they must be coerced to strings
    or the interpreter will crash.
    """
    seq = map(str, seq)
    return tuple(seq)


class Rewriter(Visitor):
    """Rewriter(obj=function or code object).
    
    Produce a new function or code object by rewriting an existing one.
    
    Notice that, unlike the base Visitor class, Rewriter does not accept a
    string or list of opcodes as an initial argument.
    """
    
    def __init__(self, obj):
        self.verbose = False
        
        if isinstance(obj, FunctionType):
            self._func = obj
            obj = obj.func_code
        
        # Copy code object attributes.
        for name in _co_code_attrs:
            value = getattr(obj, name)
            if isinstance(value, tuple):
                value = list(value)
            setattr(self, name, value)
        
        try:
            obj = obj.co_code
        except AttributeError:
            pass
        
        # Map the code block to a list of opcode numbers.
        if isinstance(obj, basestring):
            bytecode = map(ord, obj)
        elif isinstance(obj, list):
            bytecode = obj[:]
        else:
            raise TypeError("obj arg of incorrect type '%s'" % type(obj))
        
        self._bytecode = bytecode
    
    def bytecode(self):
        """Walk self and return new bytecode."""
        self.walk()
        return self.newcode
    
    def code_object(self):
        """Walk self and produce a new code object."""
        self.walk()
        codestr = ''.join(map(chr, self.newcode))
        return CodeType(self.co_argcount, self.co_nlocals, self.co_stacksize,
                        # Notice co_consts should *not* be safe_tupled.
                        self.co_flags, codestr, tuple(self.co_consts),
                        safe_tuple(self.co_names), safe_tuple(self.co_varnames),
                        self.co_filename, self.co_name, self.co_firstlineno,
                        self.co_lnotab, safe_tuple(self.co_freevars),
                        safe_tuple(self.co_cellvars))
    
    def function(self, newname=None):
        """Walk self and produce a new function."""
        try:
            f = self._func
        except AttributeError:
            if newname is None:
                newname = ''
            co = self.code_object()
            return FunctionType(co, {}, newname)
        else:
            if newname is None:
                newname = f.func_name
            co = self.code_object()
            return FunctionType(co, f.func_globals, newname,
                                f.func_defaults, f.func_closure)
    
    def const_index(self, value):
        """The index of value in co_consts, appending it if not found."""
        for pos, item in enumerate(self.co_consts):
            try:
                if type(value) == type(item) and value == item:
                    break
            except TypeError:
                pass
        else:
            pos = len(self.co_consts)
            self.co_consts.append(value)
        return pos
    
    def name_index(self, value):
        """The index of value in co_names, appending it if not found."""
        valtype = type(value)
        for pos, item in enumerate(self.co_names):
            try:
                if valtype == type(item) and value == item:
                    return pos
            except TypeError:
                pass
        
        pos = len(self.co_names)
        self.co_names.append(value)
        return pos
    
    def walk(self):
        self.newcode = []
        Visitor.walk(self)
    
    def visit_instruction(self, op, lo=None, hi=None):
        append = self.newcode.append
        append(op)
        if lo is not None:
            append(lo)
        if hi is not None:
            append(hi)
    
    def put(self, start, end, *bits):
        """Overwrite self.newcode with new opcodes (numbers or names).
        
        If the new codes are of different quantity than the old,
        modify any jump codes affected.
        """
        bitnums = numeric_opcodes(bits)
        
        # Adjust jump codes. Notice this comes before bytecode is modified.
        jca = JumpCodeAdjuster(self.newcode, start, end, len(bitnums))
        self.newcode = jca.bytecode()
        
        # Rewrite bytecode.
        self.newcode[start:end] = bitnums
    
    def tail(self, length, *bits):
        """Overwrite self.newcode[-length:] with bits."""
        end = len(self.newcode)
        self.put(end - length, end, *bits)


class Localizer(Rewriter):
    """Localizer(func, builtin_only=False, stoplist=[], verbose=False)
    
    If a global or builtin is known at compile time, replace it with a constant.
    
    This duplicates (and borrows from) Raymond Hettinger's Cookbook recipe
    at: http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/277940
    """
    def __init__(self, func, builtin_only=False, stoplist=[], verbose=False):
        Rewriter.__init__(self, func)
        
        import __builtin__
        self.env = vars(__builtin__).copy()
        if not builtin_only:
            self.env.update(func.func_globals)
        
        self.stoplist = stoplist
        self.verbose = verbose
    
    def visit_LOAD_GLOBAL(self, lo, hi):
        name = self.co_names[lo + (hi << 8)]
        if name in self.env and name not in self.stoplist:
            value = self.env[name]
            pos = self.const_index(value)
            self.tail(3, 'LOAD_CONST', pos & 0xFF, pos >> 8)
            self.debug(name, ' --> ', value)


class TaintableStack(list):
    def __init__(self, seq=[]):
        list.__init__(self, seq)
        self._taintindex = sets.Set()
        self.maxsize = len(seq)
    
    def taint(self, index=-1):
        if index < 0:
            index += len(self)
        self._taintindex.add(index)
    
    def tainted(self, index=-1):
        if index < 0:
            index = len(self) + index
        return (index in self._taintindex)
    
    def pop(self, index=-1):
        """pop(index) -> Returns a tuple!! of (value, tainted)."""
        if index < 0:
            index += len(self)
        is_tainted = (index in self._taintindex)
        if is_tainted:
            self._taintindex.remove(index)
        return list.pop(self, index), is_tainted
    
    def append(self, obj):
        list.append(self, obj)
        l = len(self)
        if l > self.maxsize:
            self.maxsize = l


class EarlyBinder(Rewriter):
    """EarlyBinder(func, reduce_getattr=True, bind_late=[])
    
    Deep-evaluate function, replacing free vars with constants.
    
    reduce_getattr: If True (the default), getattr(x, y) will be
        replaced with x.y where possible.
    
    bind_late: a list of names (globals, freevars, or attributes)
        which should not be early-bound. For example, if you want
        datetime.date.today() to be bound late, include 'today'
        in the bind_late.
    
    Example: k = lambda x: x.Date == datetime.date(2004, 1, 1)
             r = EarlyBinder(k).function()
             
    _____ k _____                _____ r _____
     0 LOAD_FAST     0 (x)        0 LOAD_FAST   0 (x)
     3 LOAD_ATTR     1 (Date)     3 LOAD_ATTR   1 (Date)
     6 LOAD_GLOBAL   2 (datetime) 6 LOAD_CONST  6 (datetime.date(2004, 1, 1))
     9 LOAD_ATTR     3 (date)
    12 LOAD_CONST    1 (2004)
    15 LOAD_CONST    2 (1)
    18 LOAD_CONST    2 (1)
    21 CALL_FUNCTION 3
    24 COMPARE_OP    2 (==)       9 COMPARE_OP  2 (==)
    27 RETURN_VALUE              12 RETURN_VALUE
    
    This also pre-computes binary operations *and all other builtin or free
    functions* where all operands are constants, globals, or freevars.
    For example:
        LOAD_CONST        1 (3)
        LOAD_CONST        2 (4)
        BINARY_MULTIPLY
        
    is replaced with:
        LOAD_CONST        5 (12)
    
    However, order is important. lambda x: x * 4 * 5 won't see any
    optimization, because the order of eval is (x * 4) * 5. Rewritten
    as lambda x: 4 * 5 * x, the "4 * 5" can be replaced with "20".
    """
    
    def __init__(self, func, reduce_getattr=True, bind_late=[]):
        Rewriter.__init__(self, func)
        self.reduce_getattr = reduce_getattr
        
        import __builtin__
        self.env = vars(__builtin__).copy()
        self.env.update(func.func_globals)
        
        # Keep a stack like the interpreter would. This does *not*