Stack and Block objects

corpora/cmps_and_ifs.py

@@ -10,4 +10,5 @@ def f3(a):
     if a <= 20:
         return True
     return False
+
 corpus = [f1, f2, f3]

func_smash.py

@@ -67,8 +67,8 @@ def print_function(func):
     i = 0
     while i < length:
         code = ord(codestring[i])
+        print opcode.opname[code].rjust(20), str(i).rjust(3),
         i += 1
-        print opcode.opname[code].rjust(20),
         if code >= opcode.HAVE_ARGUMENT:
             arg = _get_argument(codeobj, codestring, i, code)
             i += 2

prettify.py

@@ -42,139 +42,121 @@ def prettify_code(codeobj, indent=0):
         else:
             codes.append((code, None))
 
-    skip_next_line = False
-    lines = _parse_codestring(codes)
-    for i, line in enumerate(lines):
-        if skip_next_line:
-            skip_next_line = False
-            continue
-        added_indent, code = line[0], line[1:]
-        if code[0] == "else:" and lines[i+2][0] >= added_indent:
-            # Remove the automatic "pass" after each "if" as it's not needed
-            skip_next_line = True
-        _print(indent + added_indent, *code)
+    block = _parse_codestring(codes)
+    block.display(indent)
+
+def _get_func_args(func):
+    codeobj = func.__code__
+    count = codeobj.co_argcount
+    if count == 0:
+        return ""
+    return ", ".join([arg for arg in codeobj.co_varnames[:count]])
 
 def _parse_codestring(codes):
-    indent = 0
-    lines = []
-    stack = []
+    stack = Stack()
     print_buffer = []
-    block_dedent_at = []
-    block_else_at = []
+    main_block = block = Block()
+    i2block = {}
+    drops = []
+    elses = []
     i = 0
     for instruction in codes:
         code, arg = instruction
         opname = opcode.opname[code]
+        while i in elses:
+            elses.remove(i)
+            block.toggle()
+        while i in drops:
+            drops.remove(i)
+            block = block.parent()
+        i2block[i] = block
         if code >= opcode.HAVE_ARGUMENT:
             i += 3
         else:
             i += 1
-        for x in block_dedent_at:
-            if i >= x:
-                indent -= TAB
-                block_dedent_at.remove(x)
-        for x in block_else_at:
-            if i >= x:
-                lines.append((indent, "else:"))
-                indent += TAB
-                lines.append((indent, "pass"))
-                block_else_at.remove(x)
-        _print(indent, i, opname, arg, debug=True)
         if opname in OP_LOAD:
-            _push(stack, arg, literal=OP_LOAD[opname])
+            stack.push(arg, is_literal=OP_LOAD[opname])
         elif opname in OP_BINARY:
-            tos, tos1 = _pop(stack), _pop(stack)
-            _push(stack, " ".join((tos1, OP_BINARY[opname], tos)))
+            tos, tos1 = stack.pop(), stack.pop()
+            stack.push(" ".join((tos1, OP_BINARY[opname], tos)))
         elif opname in OP_INPLACE:  # Works, but doesn't use inplace op magic
-            tos, tos1 = _pop(stack), _pop(stack)
-            _push(stack, " ".join((tos1, OP_INPLACE[opname], tos)))
+            tos, tos1 = stack.pop(), stack.pop()
+            stack.push(" ".join((tos1, OP_INPLACE[opname], tos)))
         elif opname in OP_SUBSCR:
-            tos, tos1 = _pop(stack), _pop(stack)
-            _push(stack, "{0}[{1}]".format(tos1, tos))
+            tos, tos1 = stack.pop(), stack.pop()
+            stack.push("{0}[{1}]".format(tos1, tos))
         elif opname in OP_BUILD:
             args = []
             for i in xrange(arg):
-                args.append(_pop(stack))
+                args.append(stack.pop())
             args.reverse()
             start, end = OP_BUILD[opname]
-            _push(stack, start + ", ".join(args) + end)
+            stack.push(start + ", ".join(args) + end)
         elif opname == "BUILD_MAP":
-            _push(stack, "{}")
+            stack.push("{}")
         elif opname == "STORE_FAST":
-            lines.append((indent, arg, "=", _pop(stack)))
+            block.put(arg, "=", stack.pop())
         elif opname == "STORE_MAP":
-            key, value = _pop(stack), _pop(stack)
+            key, value = stack.pop(), stack.pop()
             pair = ": ".join((key, value))
-            oldmap = _pop(stack)
+            oldmap = stack.pop()
             if oldmap == "{}":
                 newmap = "{" + pair + "}"
             else:
                 newmap = oldmap[:-1] + ", " + pair + "}"
-            _push(stack, newmap)
+            stack.push(newmap)
         elif opname == "LOAD_ATTR":
-            tos = _pop(stack)
+            tos = stack.pop()
             new_tos = tos + "." + arg
-            _push(stack, new_tos)
+            stack.push(new_tos)
         elif opname == "POP_TOP":
-            lines.append((indent, _pop(stack)))
+            block.put(stack.pop())
         elif opname == "CALL_FUNCTION":
             numargs, numkwargs = arg
             args = []
             for i in xrange(numkwargs):
-                value = _pop(stack)
-                key = _pop(stack, never_literal=True)
+                value = stack.pop()
+                key = stack.pop(never_literal=True)
                 args.append("=".join((key, value)))
             for i in xrange(numargs):
-                args.append(_pop(stack))
+                args.append(stack.pop())
             args.reverse()
-            funcname = _pop(stack)
-            _push(stack, "{0}({1})".format(funcname, ", ".join(args)))
+            funcname = stack.pop()
+            stack.push("{0}({1})".format(funcname, ", ".join(args)))
         elif opname == "PRINT_ITEM":
-            print_buffer.append(_pop(stack))
+            print_buffer.append(stack.pop())
         elif opname == "PRINT_NEWLINE":
-            lines.append((indent, "print", ", ".join(print_buffer)))
+            block.put("print", ", ".join(print_buffer))
             print_buffer = []
         elif opname == "RETURN_VALUE":
-            lines.append((indent, "return", _pop(stack)))
+            block.put("return", stack.pop())
         elif opname == "COMPARE_OP":
-            tos, tos1 = _pop(stack), _pop(stack)
+            tos, tos1 = stack.pop(), stack.pop()
             compare = " ".join((tos1, arg, tos))
-            _push(stack, compare)
+            stack.push(compare)
         elif opname == "POP_JUMP_IF_FALSE":
-            test = _pop(stack)
-            block_dedent_at.append(arg)
-            block_else_at.append(arg)
-            lines.append((indent, "if {0}:".format(test)))
-            indent += TAB
-            lines.append((indent, "pass".format(test)))
+            block = block.child()
+            block.split(stack.pop())
+            elses.append(arg)
         elif opname == "POP_JUMP_IF_TRUE":
-            test = _pop(stack)
-            block_dedent_at.append(arg)
-            block_else_at.append(arg)
-            lines.append((indent, "if not ({0}):".format(test)))
-            indent += TAB
-            lines.append((indent, "pass".format(test)))
+            block = block.child()
+            block.split(stack.pop())
+            elses.append(arg)
         elif opname == "JUMP_ABSOLUTE":
-            block_dedent_at.append(i)
+            if arg < i:
+                block = block.parent()
+            else:
+                drops.append(arg)
         elif opname == "JUMP_FORWARD":
-            block_dedent_at.append(i + arg)
+            drops.append(arg + i)
+        elif opname == "SETUP_LOOP":
+            block = block.child(loop=True)
+        elif opname == "POP_BLOCK":
+            block = block.parent()
         else:
             raise NotImplementedError(opname, arg, stack)
-    return lines
 
-def _get_func_args(func):
-    codeobj = func.__code__
-    count = codeobj.co_argcount
-    if count == 0:
-        return ""
-    return ", ".join([arg for arg in codeobj.co_varnames[:count]])
-
-def _push(stack, item, literal=False):
-    stack.append((item, literal))
-
-def _pop(stack, never_literal=False):
-    item, literal = stack.pop()
-    return repr(item) if literal and not never_literal else item
+    return main_block
 
 def _print(indentation, *args, **kwargs):
     argstring = " ".join([str(arg) for arg in args if str(arg)])
@@ -185,6 +167,95 @@ def _print(indentation, *args, **kwargs):
         return
     print " " * indentation + argstring
 
+
+class Stack(object):
+    def __init__(self):
+        self._items = []
+
+    def __iter__(self):
+        while self._items:
+            yield self.pop()
+
+    def __repr__(self):
+        s = reversed([repr(itm) if lit else itm for itm, lit in self._items])
+        return "Stack[" + ", ".join(s) + "]"
+
+    def push(self, item, is_literal=False):
+        self._items.append((item, is_literal))
+
+    def pop(self, never_literal=False):
+        item, is_literal = self._items.pop()
+        return repr(item) if is_literal and not never_literal else item
+
+
+class Block(object):
+    def __init__(self, parent=None, loop=False):
+        self._parent = parent
+        self._loop = loop
+        self._true_block = self._focus = []
+        self._false_block = []
+        self._split = None
+
+    def __iter__(self):
+        for item in self._render_lines():
+            yield item[1]
+
+    def __repr__(self):
+        return str([item[1] for item in self._render_lines()])
+
+    def _has_lines(self):
+        return self._split or self._true_block or self._false_block
+
+    def _render_subblock(self, block, lines, indent):
+        for item in block:
+            if isinstance(item, Block):
+                lines += item._render_lines(indent)
+            else:
+                lines.append((indent, item))
+
+    def _render_lines(self, indent=0):
+        lines = []
+        if self._split:
+            lines.append((indent, (self._split,)))
+            indent += TAB
+        self._render_subblock(self._true_block, lines, indent)
+        if self._false_block:
+            lines.append((indent - TAB, ("else:",)))
+            self._render_subblock(self._false_block, lines, indent)
+        return lines
+
+    def child(self, loop=False):
+        if self._loop and not loop and not self._has_lines():
+            loop = True 
+        child = Block(parent=self, loop=loop)
+        self._focus.append(child)
+        return child
+
+    def parent(self):
+        if self._parent:
+            return self._parent
+        raise RuntimeError("Popping an orphaned block")
+
+    def put(self, *code):
+        self._focus.append(code)
+
+    def split(self, test):
+        if self._loop:
+            self._split = "while " + test + ":"
+        else:
+            self._split = "if " + test + ":"
+
+    def toggle(self):
+        if self._focus is self._true_block:
+            self._focus = self._false_block
+        else:
+            self._focus = self._true_block
+
+    def display(self, indent=0):
+        for indent, code in self._render_lines(indent):
+            _print(indent, *code)
+
+
 if __name__ == "__main__":
     def f1(a):
         b = a + 10 / 7.0
@@ -221,6 +292,23 @@ if __name__ == "__main__":
                             line10
         return line11
 
+    def f3(x, y):
+        if cmp1:
+            while cmp1:
+                if cmp2:
+                    while cmp3:
+                        line1
+        else:
+            while cmp4:
+                line2
+                line3
+                line4
+                if cmp5:
+                    line5
+        return line6
+
     prettify_function(f1)
     print
     prettify_function(f2)
+    print
+    prettify_function(f3)