# HG changeset patch
# User Thinker K.F. Li <thinker@codemud.net>
# Date 1312180624 -28800
# Node ID 835336632aba8422eebd62e2af587f8efa52562b
# Parent  18be67af7f1e3f0c62b3de19d9d2ab5cc721ab69
Add collect_typeidxs_in_method()

diff -r 18be67af7f1e -r 835336632aba paraspace/dexfile.py
--- a/paraspace/dexfile.py	Mon Aug 01 12:27:28 2011 +0800
+++ b/paraspace/dexfile.py	Mon Aug 01 14:37:04 2011 +0800
@@ -1902,6 +1902,19 @@
                    for tl_typeid in tl_typeids]
         return typeids
 
+    ## \brief Return code block of given method.
+    #
+    # Code block is a string of byte code instructions for Dalvik VM.
+    #
+    @staticmethod
+    def get_code_block_method(method):
+        if not method.codeOffRef.is_true:
+            return ''
+
+        code = method.codeOffRef.value
+        insns = code.insns.data
+        return insns
+
     ## \brief Dump content of a proto ID.
     @staticmethod
     def dump_protoid(protoid):
diff -r 18be67af7f1e -r 835336632aba paraspace/injection.py
--- a/paraspace/injection.py	Mon Aug 01 12:27:28 2011 +0800
+++ b/paraspace/injection.py	Mon Aug 01 14:37:04 2011 +0800
@@ -364,3 +364,55 @@
         class_redirect_types(dex, classdef, types_redir, methods_redir)
         pass
     pass
+
+
+## \brief Collect all type indices mentioned in the code of given method.
+#
+# \param method is a \ref _DEX_Method.
+#
+def collect_typeidxs_in_method(dex, method):
+    from paraspace.dexfile import _DEX_Method, DEXFile_linked
+    from paraspace.dalvik_opcodes import decode_insn_blk, all_opcodes
+    from itertools import chain
+    
+    assert isinstance(method, _DEX_Method)
+
+    def get_typeidx_methodidx(methodidx):
+        methodid = dex.find_methodid_idx(methodidx)
+        method_typeid = methodid.classIdx
+        method_typeidx = dex.get_idx_typeid(method_typeid)
+        return method_typeidx
+
+    def collect_types_in_op_vector(op_vector):
+        code, args = op_vector
+        
+        if code == all_opcodes.OP_NEW_INSTANCE:
+            return (args[1],)
+        
+        if code in (all_opcodes.OP_INVOKE_DIRECT,
+                    all_opcodes.OP_INVOKE_VIRTUAL,
+                    all_opcodes.OP_INVOKE_SUPER,
+                    all_opcodes.OP_INVOKE_STATIC,
+                    all_opcodes.OP_INVOKE_INTERFACE):
+            methodidx = args[2]
+            method_typeidx = get_typeidx_methodidx(methodidx)
+            return (method_typeidx,)
+
+        if code in (all_opcodes.OP_INVOKE_VIRTUAL_RANGE,
+                    all_opcodes.OP_INVOKE_DIRECT_RANGE,
+                    all_opcodes.OP_INVOKE_SUPER_RANGE,
+                    all_opcodes.OP_INVOKE_STATIC_RANGE,
+                    all_opcodes.OP_INVOKE_INTERFACE_RANGE):
+            methodidx = args[1]
+            method_typeidx = get_typeidx_methodidx(methodidx)
+            return (method_typeidx,)
+        
+        return ()
+
+    code_blk = DEXFile_linked.get_code_block_method(method)
+    op_vectors = decode_insn_blk(code_blk)
+    types_insns = [collect_types_in_op_vector(op_vector)
+                   for op_vector in op_vectors]
+    typeidxs = list(chain(*types_insns))
+    
+    return typeidxs
diff -r 18be67af7f1e -r 835336632aba paraspace/tests/injection_test.py
--- a/paraspace/tests/injection_test.py	Mon Aug 01 12:27:28 2011 +0800
+++ b/paraspace/tests/injection_test.py	Mon Aug 01 14:37:04 2011 +0800
@@ -172,3 +172,30 @@
             pass
         pass
     pass
+
+
+def collect_types_in_method_test():
+    from paraspace.dex_deptracker import prepare_dep_decls
+    from paraspace.injection import collect_typeidxs_in_method
+    
+    _install_dexfile_4_deptracker()
+    
+    all_dep_decls = prepare_dep_decls()
+    
+    srcdir = os.path.dirname(__file__)
+    srcroot = os.path.join(srcdir, '..', '..')
+    
+    fakefile_fn = os.path.join(srcroot, 'data', 'fakefile.dex')
+    fakefile_dex = dexfile.DEXFile.open(fakefile_fn)
+    fakefile_linked = \
+        dexfile.DEXFile_linked. \
+        build_dependencies(fakefile_dex, all_dep_decls)
+    
+    fakefile_def = \
+        fakefile_linked.find_class_name('Lcom/codemud/fakefile/fakefile;')
+
+    init_method = fakefile_linked.find_method_name('<init>', fakefile_def)
+
+    typeidxs = collect_typeidxs_in_method(fakefile_linked, init_method)
+    assert len(typeidxs) == 1
+    pass