Improve ir hashability and type correctness

Author: twizmwazin

pysoot/sootir/soot_class.py

@@ -2,6 +2,8 @@
 
 from dataclasses import dataclass
 
+from frozendict import frozendict
+
 from .soot_method import SootMethod
 from . import convert_soot_attributes
 
@@ -18,10 +20,10 @@ class SootClass:
     ]  # TODO: replace with dataclass in Python 3.10
     name: str
     super_class: str
-    interfaces: list[str]
-    attrs: list[str]
-    methods: list[SootMethod]
-    fields: dict[str, tuple[list[str], str]]
+    interfaces: tuple[str, ...]
+    attrs: tuple[str, ...]
+    methods: tuple[SootMethod, ...]
+    fields: frozendict[str, tuple[tuple[str], str]]
 
     def __str__(self):
         tstr = "//" + repr(self) + "\n"
@@ -73,7 +75,7 @@ def from_ir(ir_class):
         fields = {}
         for field in ir_class.getFields():
             fields[str(field.getName())] = (
-                convert_soot_attributes(field.getModifiers()),
+                tuple(convert_soot_attributes(field.getModifiers())),
                 str(field.getType()),
             )
 
@@ -83,5 +85,5 @@ def from_ir(ir_class):
         else:
             super_class = ""
         return SootClass(
-            class_name, super_class, interface_names, attrs, methods, fields
+            class_name, super_class, tuple(interface_names), tuple(attrs), tuple(methods), frozendict(fields)
         )

pysoot/sootir/soot_expr.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any
 
 from .soot_value import SootValue
 
@@ -152,7 +151,7 @@ def from_ir(type_, expr_name, ir_subvalue):
 class SootNewMultiArrayExpr(SootExpr):
     __slots__ = ["base_type", "sizes"]  # TODO: replace with dataclass in Python 3.10
     base_type: str
-    sizes: Any
+    sizes: tuple[SootValue, ...]
 
     def __str__(self):
         return "new %s%s" % (
@@ -165,7 +164,7 @@ def from_ir(type_, expr_name, ir_subvalue):
         return SootNewMultiArrayExpr(
             type_,
             str(ir_subvalue.getBaseType()),
-            (SootValue.from_ir(size) for size in ir_subvalue.getSizes()),
+            tuple(SootValue.from_ir(size) for size in ir_subvalue.getSizes()),
         )
 
 
@@ -195,7 +194,7 @@ def __str__(self):
     @staticmethod
     def from_ir(type_, expr_name, ir_subvalue):
         return SootPhiExpr(
-            type_, (SootValue.from_ir(v.getValue()) for v in ir_subvalue.getArgs())
+            type_, tuple(SootValue.from_ir(v.getValue()) for v in ir_subvalue.getArgs())
         )
 
 
@@ -211,8 +210,8 @@ class SootInvokeExpr(SootExpr):
     ]  # TODO: replace with dataclass in Python 3.10
     class_name: str
     method_name: str
-    method_params: Any
-    args: Any
+    method_params: tuple[str, ...]
+    args: tuple[SootValue, ...]
 
     def __str__(self):
         return "%s.%s(%s)]" % (
@@ -229,7 +228,7 @@ def list_to_arg_str(args):
 @dataclass(unsafe_hash=True)
 class SootVirtualInvokeExpr(SootInvokeExpr):
     __slots__ = ["base"]  # TODO: replace with dataclass in Python 3.10
-    base: Any
+    base: SootValue
 
     def __str__(self):
         return "%s.%s(%s) [virtualinvoke %s" % (
@@ -261,8 +260,9 @@ class SootDynamicInvokeExpr(SootInvokeExpr):
         "bootstrap_method",
         "bootstrap_args",
     ]  # TODO: replace with dataclass in Python 3.10
-    bootstrap_method: Any
-    bootstrap_args: Any
+    # TODO:  bootstrap_method and bootstrap_args are not implemented yet
+    bootstrap_method: None
+    bootstrap_args: None
 
     @staticmethod
     def from_ir(type_, expr_name, ir_expr):
@@ -280,7 +280,7 @@ def from_ir(type_, expr_name, ir_expr):
             class_name=class_name,
             method_name=method_name,
             method_params=method_params,
-            args=method_args,
+            args=args,
             bootstrap_method=bootstrap_method,
             bootstrap_args=bootstrap_args,
         )
@@ -289,7 +289,7 @@ def from_ir(type_, expr_name, ir_expr):
 @dataclass(unsafe_hash=True)
 class SootInterfaceInvokeExpr(SootInvokeExpr):
     __slots__ = ["base"]  # TODO: replace with dataclass in Python 3.10
-    base: Any
+    base: SootValue
 
     def __str__(self):
         return "%s.%s(%s) [interfaceinvoke %s" % (
@@ -303,7 +303,7 @@ def __str__(self):
     def from_ir(type_, expr_name, ir_expr):
         args = tuple([SootValue.from_ir(arg) for arg in ir_expr.getArgs()])
         called_method = ir_expr.getMethod()
-        params = tuple([str(param) for param in called_method.getParameterTypes()])
+        params = tuple(str(param) for param in called_method.getParameterTypes())
 
         return SootInterfaceInvokeExpr(
             type=type_,
@@ -318,7 +318,7 @@ def from_ir(type_, expr_name, ir_expr):
 @dataclass(unsafe_hash=True)
 class SootSpecialInvokeExpr(SootInvokeExpr):
     __slots__ = ["base"]  # TODO: replace with dataclass in Python 3.10
-    base: Any
+    base: SootValue
 
     def __str__(self):
         return "%s.%s(%s) [specialinvoke %s" % (
@@ -330,9 +330,9 @@ def __str__(self):
 
     @staticmethod
     def from_ir(type_, expr_name, ir_expr):
-        args = tuple([SootValue.from_ir(arg) for arg in ir_expr.getArgs()])
+        args = tuple(SootValue.from_ir(arg) for arg in ir_expr.getArgs())
         called_method = ir_expr.getMethod()
-        params = tuple([str(param) for param in called_method.getParameterTypes()])
+        params = tuple(str(param) for param in called_method.getParameterTypes())
 
         return SootSpecialInvokeExpr(
             type=type_,
@@ -357,9 +357,9 @@ def __str__(self):
 
     @staticmethod
     def from_ir(type_, expr_name, ir_expr):
-        args = tuple([SootValue.from_ir(arg) for arg in ir_expr.getArgs()])
+        args = tuple(SootValue.from_ir(arg) for arg in ir_expr.getArgs())
         called_method = ir_expr.getMethod()
-        params = tuple([str(param) for param in called_method.getParameterTypes()])
+        params = tuple(str(param) for param in called_method.getParameterTypes())
 
         return SootStaticInvokeExpr(
             type=type_,

pysoot/sootir/soot_method.py

@@ -4,6 +4,7 @@
 from dataclasses import dataclass
 from functools import lru_cache
 
+from frozendict import frozendict
 from jpype.types import JClass
 
 from .soot_block import SootBlock
@@ -31,11 +32,11 @@ class SootMethod:
     exceptions: tuple[str, ...]
     blocks: tuple[SootBlock, ...]
     params: tuple[str, ...]
-    basic_cfg: defaultdict[SootBlock, list[SootBlock]]
-    exceptional_preds: defaultdict[SootBlock, list[SootBlock]]
+    basic_cfg: frozendict[SootBlock, tuple[SootBlock]]
+    exceptional_preds: frozendict[SootBlock, tuple[SootBlock]]
 
     @property
-    # @lru_cache(maxsize=1)
+    @lru_cache(maxsize=1)
     def block_by_label(self):
         return {b.label: b for b in self.blocks}
 
@@ -109,13 +110,13 @@ def from_ir(class_name, ir_method):
                     if "Assign" in ir_stmt.getClass().getSimpleName():
                         ir_expr = ir_stmt.getRightOp()
                         if "Phi" in ir_expr.getClass().getSimpleName():
-                            values = [
+                            values = tuple(
                                 (
                                     SootValue.from_ir(v.getValue()),
                                     stmt_to_block_idx[v.getUnit()],
                                 )
                                 for v in ir_expr.getArgs()
-                            ]
+                            )
 
                             phi_expr = SootValue.IREXPR_TO_EXPR[ir_expr]
                             phi_expr.values = values
@@ -136,6 +137,6 @@ def from_ir(class_name, ir_method):
             attrs=tuple(attrs),
             exceptions=tuple(exceptions),
             blocks=tuple(blocks),
-            basic_cfg=basic_cfg,
-            exceptional_preds=exceptional_preds,
+            basic_cfg=frozendict({k: tuple(v) for k, v in basic_cfg.items()}),
+            exceptional_preds=frozendict({k: tuple(v) for k, v in exceptional_preds.items()}),
         )

pysoot/sootir/soot_statement.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any
 
 from frozendict import frozendict
 
@@ -197,7 +196,7 @@ class LookupSwitchStmt(SootStmt):
         "default_target",
     ]  # TODO: replace with dataclass in Python 3.10
     key: SootValue
-    lookup_values_and_targets: Any
+    lookup_values_and_targets: frozendict[int, SootStmt]
     default_target: SootStmt
 
     def __str__(self):
@@ -214,11 +213,11 @@ def from_ir(label, offset, ir_stmt, stmt_map=None):
         lookup_values_and_targets = frozendict({k: v for k, v in zip(lookup_values, targets)})
 
         return LookupSwitchStmt(
-            label,
-            offset,
-            SootValue.from_ir(ir_stmt.getKey()),
-            lookup_values_and_targets,
-            stmt_map[ir_stmt.getDefaultTarget()],
+            label=label,
+            offset=offset,
+            key=SootValue.from_ir(ir_stmt.getKey()),
+            lookup_values_and_targets=lookup_values_and_targets,
+            default_target=stmt_map[ir_stmt.getDefaultTarget()],
         )
 
 
@@ -233,10 +232,10 @@ class TableSwitchStmt(SootStmt):
         "default_target",
     ]
     key: SootValue
-    low_index: Any
-    high_index: Any
-    targets: Any
-    lookup_values_and_targets: Any
+    low_index: int
+    high_index: int
+    targets: tuple[SootStmt, ...]
+    lookup_values_and_targets: frozendict[int, SootStmt]
     default_target: SootStmt
 
     def __str__(self):
@@ -248,21 +247,21 @@ def __str__(self):
 
     @staticmethod
     def from_ir(label, offset, ir_stmt, stmt_map=None):
-        targets = [stmt_map[t] for t in ir_stmt.getTargets()]
+        targets = tuple(stmt_map[t] for t in ir_stmt.getTargets())
         dict_iter = zip(
             range(ir_stmt.getLowIndex(), ir_stmt.getHighIndex() + 1), targets
         )
         lookup_values_and_targets = {k: v for k, v in dict_iter}
 
         return TableSwitchStmt(
-            label,
-            offset,
-            SootValue.from_ir(ir_stmt.getKey()),
-            ir_stmt.getLowIndex(),
-            ir_stmt.getHighIndex(),
-            tuple(targets),
-            stmt_map[ir_stmt.getDefaultTarget()],
-            frozendict(lookup_values_and_targets),
+            label=label,
+            offset=offset,
+            key=SootValue.from_ir(ir_stmt.getKey()),
+            low_index=int(ir_stmt.getLowIndex()),
+            high_index=int(ir_stmt.getHighIndex()),
+            targets=tuple(targets),
+            default_target=stmt_map[ir_stmt.getDefaultTarget()],
+            lookup_values_and_targets=frozendict(lookup_values_and_targets),
         )
 
 

pysoot/sootir/soot_value.py

@@ -1,9 +1,6 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Any
-
-from jpype.types import JDouble, JFloat, JInt, JLong, JString
 
 
 @dataclass(unsafe_hash=True)
@@ -53,8 +50,8 @@ def from_ir(type_, ir_value):
 @dataclass(unsafe_hash=True)
 class SootArrayRef(SootValue):
     __slots__ = ["base", "index"]  # TODO: replace with dataclass in Python 3.10
-    base: Any
-    index: Any
+    base: SootValue
+    index: SootValue
 
     def __str__(self):
         return "%s[%s]" % (self.base, self.index)
@@ -83,14 +80,14 @@ def from_ir(type_, ir_value):
 @dataclass(unsafe_hash=True)
 class SootParamRef(SootValue):
     __slots__ = ["index"]  # TODO: replace with dataclass in Python 3.10
-    index: Any
+    index: int
 
     def __str__(self):
         return "@parameter%d[%s]" % (self.index, self.type)
 
     @staticmethod
     def from_ir(type_, ir_value):
-        return SootParamRef(type_, ir_value.getIndex())
+        return SootParamRef(type_, int(ir_value.getIndex()))
 
 
 @dataclass(unsafe_hash=True)
@@ -108,7 +105,7 @@ def from_ir(type_, ir_value):
 @dataclass(unsafe_hash=True)
 class SootStaticFieldRef(SootValue):
     __slots__ = ["field"]  # TODO: replace with dataclass in Python 3.10
-    field: Any
+    field: tuple[str, str]
 
     def __str__(self):
         return "StaticFieldRef %s" % (self.field,)
@@ -122,7 +119,7 @@ def from_ir(type_, ir_value):
 @dataclass(unsafe_hash=True)
 class SootInstanceFieldRef(SootValue):
     __slots__ = ["base", "field"]  # TODO: replace with dataclass in Python 3.10
-    base: Any
+    base: SootValue
     field: tuple[str, str]
 
     def __str__(self):
@@ -141,14 +138,14 @@ def from_ir(type_, ir_value):
 @dataclass(unsafe_hash=True)
 class SootClassConstant(SootValue):
     __slots__ = ["value"]  # TODO: replace with dataclass in Python 3.10
-    value: Any
+    value: str
 
     def __str__(self):
         return str(self.value)
 
     @staticmethod
     def from_ir(type_, ir_value):
-        return SootClassConstant(type_, ir_value)
+        return SootClassConstant(type_, str(ir_value.getValue()))
 
 
 @dataclass(unsafe_hash=True)

tests/test_pysoot.py

@@ -81,7 +81,7 @@ def test_exceptions1(self):
         lifter = Lifter(jar)
 
         mm = lifter.classes["exceptions1.Main"].methods[1]
-        assert mm.basic_cfg[mm.blocks[0]] == [mm.blocks[1], mm.blocks[2]]
+        assert mm.basic_cfg[mm.blocks[0]] == (mm.blocks[1], mm.blocks[2])
         assert len(mm.exceptional_preds) == 1
 
         preds = mm.exceptional_preds[mm.blocks[18]]