collect data from instructions

androscalpel/src/code.rs

@@ -14,7 +14,7 @@ use crate::{DexString, IdField, IdMethod, IdMethodType, IdType, Instruction, Met
 #[pyclass]
 #[derive(Debug, Clone)]
 pub struct Code {
-    // TODO: remove and compute this value from code.
+    // TODO: remove and compute this value from code?
     /// The number of registers used by the code
     #[pyo3(get, set)]
     pub registers_size: u16,
@@ -30,27 +30,28 @@ pub struct Code {
     /// The debug info
     #[pyo3(get, set)]
     pub debug_info: Vec<u8>,
-    // TODO: implement OPcode
     /// The instructions.
     #[pyo3(get, set)]
     pub insns: Vec<Instruction>,
-    // TODO: currently unusable, juste a mapping ty TryItem
-    // TODO: maybe implement as custom OPcode to make me easy to modify?
-    // /// Try blocks
-    // #[pyo3(get, set)]
-    // pub tries: Vec<(u32, u16, TmpHandlerType)>,
-    // TODO: currently unusable, juste a mapping ty TryItem
-    // TODO: maybe implement as custom OPcode to make me easy to modify?
 }
 
 #[pymethods]
 impl Code {
-    /*
     #[new]
-    pub fn new() -> Self {
-        todo!()
+    pub fn new(
+        registers_size: u16,
+        ins_size: u16,
+        outs_size: u16,
+        insns: Vec<Instruction>,
+    ) -> Self {
+        Self {
+            registers_size,
+            ins_size,
+            outs_size,
+            insns,
+            debug_info: vec![],
+        }
     }
-    */
 
     pub fn __str__(&self) -> String {
         self.__repr__()
@@ -62,53 +63,57 @@ impl Code {
 
     /// Return all strings referenced in the codes.
     pub fn get_all_strings(&self) -> HashSet<DexString> {
-        todo!()
-        /*
         let mut strings = HashSet::new();
-        for (_, _, (list, _)) in &self.tries {
-            for (ty, _) in list {
-                strings.extend(ty.get_all_strings());
-            }
+        // TODO debug info contains types and strings
+        for ins in &self.insns {
+            strings.extend(ins.get_all_strings());
         }
         strings
-        */
     }
 
     /// Return all types referenced in the codes.
     pub fn get_all_types(&self) -> HashSet<IdType> {
-        todo!()
-        /*
-        let mut types = HashSet::new();
-        for (_, _, (list, _)) in &self.tries {
-            for (ty, _) in list {
-                types.insert(ty.clone());
-            }
+        let mut type_ids = HashSet::new();
+        // TODO debug info contains types and strings
+        for ins in &self.insns {
+            type_ids.extend(ins.get_all_types());
         }
-        types
-        */
+        type_ids
     }
 
     /// Return all prototypes referenced in the codes.
     pub fn get_all_protos(&self) -> HashSet<IdMethodType> {
-        // TODO
-        HashSet::new()
+        let mut protos = HashSet::new();
+        for ins in &self.insns {
+            protos.extend(ins.get_all_protos());
+        }
+        protos
     }
 
     /// Return all field ids referenced in the codes.
     pub fn get_all_field_ids(&self) -> HashSet<IdField> {
-        // TODO
-        HashSet::new()
+        let mut fields = HashSet::new();
+        for ins in &self.insns {
+            fields.extend(ins.get_all_field_ids());
+        }
+        fields
     }
 
     /// Return all method ids referenced in the codes.
     pub fn get_all_method_ids(&self) -> HashSet<IdMethod> {
-        // TODO
-        HashSet::new()
+        let mut methods = HashSet::new();
+        for ins in &self.insns {
+            methods.extend(ins.get_all_method_ids())
+        }
+        methods
     }
 
     /// Return all method handles referenced in the codes.
     pub fn get_all_method_handles(&self) -> HashSet<MethodHandle> {
-        // TODO
-        HashSet::new()
+        let mut handles = HashSet::new();
+        for ins in &self.insns {
+            handles.extend(ins.get_all_method_handles());
+        }
+        handles
     }
 }