From 75aaa1f1d0e62a7ef9a5ed901bb21ba8b3b93471 Mon Sep 17 00:00:00 2001
From: Ben Kurtovic <ben.kurtovic@verizon.net>
Date: Mon, 23 Apr 2012 21:47:46 -0400
Subject: [PATCH] Some more corpora; updates

---
 corpora/functions.py |  2 +-
 corpora/loops.py     |  5 +++++
 corpora/objects.py   | 13 +++++++++++++
 func_smash.py        | 21 ++++++++-------------
 4 files changed, 27 insertions(+), 14 deletions(-)
 create mode 100644 corpora/loops.py
 create mode 100644 corpora/objects.py

diff --git a/corpora/functions.py b/corpora/functions.py
index 301dd0b..135c1dd 100644
--- a/corpora/functions.py
+++ b/corpora/functions.py
@@ -1,5 +1,5 @@
 def f1(a):
-    return str(a)
+    return int(a)
 
 def f2(a):
     return (lambda i: len(str(i)))(a)
diff --git a/corpora/loops.py b/corpora/loops.py
new file mode 100644
index 0000000..7bd7d53
--- /dev/null
+++ b/corpora/loops.py
@@ -0,0 +1,5 @@
+def f1(a):
+    for i in a:
+        print i
+
+corpus = [f1]
diff --git a/corpora/objects.py b/corpora/objects.py
new file mode 100644
index 0000000..c3a231b
--- /dev/null
+++ b/corpora/objects.py
@@ -0,0 +1,13 @@
+def f1(a):
+    return [a] * 4
+
+def f2(a):
+    return (a,) + (1, 2)
+
+def f3(a):
+    return {a: 3}
+
+def f4(a):
+    return {a, 5}
+
+corpus = [f1, f2, f3, f4]
diff --git a/func_smash.py b/func_smash.py
index 109ef79..f1961ed 100644
--- a/func_smash.py
+++ b/func_smash.py
@@ -8,9 +8,12 @@ import sys
 import types
 
 OPMAP = opcode.opmap
+OP_HASBUILD = [OPMAP[n] for n in ("BUILD_TUPLE", "BUILD_LIST", "BUILD_MAP",
+                                  "BUILD_SET")]
 OP_HASCALL = [OPMAP[n] for n in ("CALL_FUNCTION", "CALL_FUNCTION_VAR",
                                  "CALL_FUNCTION_KW", "CALL_FUNCTION_VAR_KW")]
 OP_HASMAKE = [OPMAP[n] for n in ("MAKE_FUNCTION", "MAKE_CLOSURE")]
+OP_LITERALARG = opcode.hasjabs + opcode.hasjrel + OP_HASBUILD + OP_HASMAKE
 
 MARKOV_START = -1
 MARKOV_END = -2
@@ -73,6 +76,8 @@ def print_function(func):
                 print " (to {0})".format(arg)
             elif code in opcode.hasjrel:
                 print " (+{0})".format(arg)
+            elif code in OP_HASBUILD:
+                print " ({0} items)".format(arg)
             elif code in OP_HASCALL:
                 print " ({0} args, {1} kwargs)".format(*arg)
             elif code in OP_HASMAKE:
@@ -98,6 +103,7 @@ def run():
 
 def _parse_func(func, chain):
     codeobj = func.__code__
+    print codeobj.co_argcount
     codestring = codeobj.co_code
     length = len(codestring)
     i = 0
@@ -124,13 +130,9 @@ def _get_argument(codeobj, codestring, i, code):
         return codeobj.co_varnames[arg]
     elif code in opcode.hascompare:
         return opcode.cmp_op[arg]
-    elif code in opcode.hasjabs:
-        return arg
-    elif code in opcode.hasjrel:
-        return arg - i - 2
     elif code in OP_HASCALL:
         return (ord(codestring[i]), ord(codestring[i + 1]))
-    elif code in OP_HASMAKE:
+    elif code in OP_LITERALARG:
         return arg
     raise NotImplementedError(code, opcode.opname[code])
 
@@ -144,15 +146,12 @@ def _make_codes(chain):
     codes = []
     instruction = random.choice(chain[MARKOV_START])
     constants, names, varnames = [], [], []
-    i = 0
     while 1:
-        i += 1
         code, arg = instruction
         if code == MARKOV_END:
             break
         codes.append(code)
         if code >= opcode.HAVE_ARGUMENT:
-            i += 2
             if code in opcode.hasconst:
                 if arg not in constants:
                     constants.append(arg)
@@ -167,14 +166,10 @@ def _make_codes(chain):
                 _coerce_arg_into_codes(codes, varnames.index(arg))
             elif code in opcode.hascompare:
                 _coerce_arg_into_codes(codes, opcode.cmp_op.index(arg))
-            elif code in opcode.hasjabs:
-                _coerce_arg_into_codes(codes, arg)
-            elif code in opcode.hasjrel:
-                _coerce_arg_into_codes(codes, arg + i)
             elif code in OP_HASCALL:
                 codes.append(arg[0])
                 codes.append(arg[1])
-            elif code in OP_HASMAKE:
+            elif code in OP_LITERALARG:
                 _coerce_arg_into_codes(codes, arg)
             else:
                 raise NotImplementedError(code, opcode.opname[code])