--- a/paraspace/injection.py	Mon Aug 15 21:51:39 2011 +0800
+++ b/paraspace/injection.py	Tue Aug 16 14:53:41 2011 +0800
@@ -268,6 +268,30 @@
     return clone
 
 
+## \biref Relative object set for a _DEX_ClassData.
+class classdef_rel_set(object):
+    classdefs = None
+    typeids = None
+    strids = None
+    
+    def __init__(self, classdefs, typeids, strids):
+        self.classdefs = set(classdefs)
+        self.typeids = set(typeids)
+        self.strids = set(strids)
+        pass
+
+    @staticmethod
+    def dup(src):
+        return relative_obj_set(src.classdefs, src.typeids, src.strids)
+
+    def update(self, new_data):
+        self.classdefs.update(new_data.classdefs)
+        self.typeids.update(new_data.typeids)
+        self.strids.update(new_data.strids)
+        pass
+    pass
+
+
 ## \brief Collect info of classes mentioned by the code of given class.
 def _find_class_relative(dex, classdef):
     def classify_typeids_defined(dex, typeids):
@@ -292,7 +316,8 @@
     
     classdefs, typeids = classify_typeids_defined(dex, typeids)
     
-    return classdefs, typeids, strids
+    rel_set = classdef_rel_set(classdefs, typeids, strids)
+    return rel_set
 
 
 def dexfile_insert_classdefs(dex_dst, dex_src, classdefs):
@@ -369,22 +394,18 @@
 
 ## \brief Collects relative type IDs and classes definition for given class.
 def collect_classdefs_relative(dex, classdefs):
-    rel_classdefs = set(classdefs)
-    rel_typeids = set()
-    rel_strids = set()
+    rel_set = classdef_rel_set(classdefs, [], [])
     
     classdef_queue = list(classdefs)
     while classdef_queue:
         cur_classdef = classdef_queue.pop(0)
         
-        classdefs, typeids, strids = _find_class_relative(dex, cur_classdef)
-        rel_typeids.update(typeids)
-        new_classdefs = list(set(classdefs) - rel_classdefs)
-        classdef_queue = classdef_queue + new_classdefs
-        rel_classdefs.update(new_classdefs)
-        rel_strids.update(strids)
+        new_rel_set = _find_class_relative(dex, cur_classdef)
+        rel_set.update(new_rel_set)
+        new_classdefs = new_rel_set.classdefs - rel_set.classdefs
+        classdef_queue = classdef_queue + list(new_classdefs)
         pass
-    return rel_classdefs, rel_typeids, rel_strids
+    return rel_set
 
 
 ## \brief Clone and insert given and relative classes into another DEXFile.
@@ -405,23 +426,24 @@
             return True
         return False
 
-    relative_classdefs, relative_typeids, relative_strids = \
-        collect_classdefs_relative(dex_src, classdefs)
+    rel_set = collect_classdefs_relative(dex_src, classdefs)
     
-    for classdef in relative_classdefs:
+    for classdef in rel_set.classdefs:
         if classdef_not_in_dst(classdef):
             continue
         raise ValueError, '%s is already in DEX %s: can not insert it' % \
             (repr(classdef), repr(dex_dst))
     
     cloning_classdefs = \
-        dexfile_insert_classdefs(dex_dst, dex_src, relative_classdefs)
+        dexfile_insert_classdefs(dex_dst, dex_src, rel_set.classdefs)
     cloning_typeids = \
-        dexfile_insert_or_merge_typeids(dex_dst, dex_src, relative_typeids)
+        dexfile_insert_or_merge_typeids(dex_dst, dex_src, rel_set.typeids)
     cloning_strids = \
-        dexfile_insert_stringids(dex_dst, dex_src, relative_strids)
+        dexfile_insert_stringids(dex_dst, dex_src, rel_set.strids)
 
-    return cloning_classdefs, cloning_typeids, cloning_strids
+    cloning = classdef_rel_set(cloning_classdefs, cloning_typeids,
+                               cloning_strids)
+    return cloning
 
 
 ## \brief Redirect types, methods and strings for the code of given method.
@@ -724,13 +746,12 @@
 
 ## \brief Redirect code for methods of injected classes.
 def redirect_injected_code(dex_dst, dex_src, classdefs):
-    relative_classdefs, relative_typeids, relative_strids = \
-        collect_classdefs_relative(dex_src, classdefs)
+    rel_set = collect_classdefs_relative(dex_src, classdefs)
     
     typeidxs_redir = \
         make_typeidxs_map_after_injection(dex_dst, dex_src, \
-                                              relative_classdefs, \
-                                              relative_typeids)
+                                              rel_set.classdefs, \
+                                              rel_set.typeids)
     methodidxs_redir = \
         make_methodidxs_redir_map(dex_src, dex_dst, typeidxs_redir)
 
@@ -758,10 +779,9 @@
     assert isinstance(dex_dst, DEXFile_linked)
     assert isinstance(dex_src, DEXFile_linked)
 
-    injected_classdefs, injected_typeids, injected_strids = \
-        dexfile_insert_classdefs_relative(dex_dst, dex_src, classdefs)
+    rel_set = dexfile_insert_classdefs_relative(dex_dst, dex_src, classdefs)
     redirect_injected_code(dex_dst, dex_src, classdefs)
-    return injected_classdefs
+    return list(rel_set.classdefs)
 
 
 ## \brief Redirect all references of given class and its methods.

--- a/paraspace/tests/injection_test.py	Mon Aug 15 21:51:39 2011 +0800
+++ b/paraspace/tests/injection_test.py	Tue Aug 16 14:53:41 2011 +0800
@@ -263,10 +263,10 @@
     assert fakefile_dataheader.directMethodsSize == 1
     assert fakefile_dataheader.virtualMethodsSize == 0
 
-    classdefs, typeids, strids = \
-        dexfile_insert_classdefs_relative(helloworld_linked, \
-                                              fakefile_linked, \
-                                              [fakefile_def])
+    rel_set = dexfile_insert_classdefs_relative(helloworld_linked,
+                                                fakefile_linked,
+                                                [fakefile_def])
+    classdefs = list(rel_set.classdefs)
     assert classdefs
     assert len(classdefs) == 1
     assert classdefs[0] != fakefile_def