Merge pull request #1757 from webthethird/dev-upgradeability-utils

Add upgradeability utils

slither/core/declarations/contract.py

@@ -93,6 +93,9 @@ def __init__(self, compilation_unit: "SlitherCompilationUnit", scope: "FileScope
         self._signatures: Optional[List[str]] = None
         self._signatures_declared: Optional[List[str]] = None
 
+        self._fallback_function: Optional["FunctionContract"] = None
+        self._receive_function: Optional["FunctionContract"] = None
+
         self._is_upgradeable: Optional[bool] = None
         self._is_upgradeable_proxy: Optional[bool] = None
         self._upgradeable_version: Optional[str] = None
@@ -663,6 +666,24 @@ def functions_and_modifiers_declared(self) -> List["Function"]:
         """
         return self.functions_declared + self.modifiers_declared  # type: ignore
 
+    @property
+    def fallback_function(self) -> Optional["FunctionContract"]:
+        if self._fallback_function is None:
+            for f in self.functions:
+                if f.is_fallback:
+                    self._fallback_function = f
+                    break
+        return self._fallback_function
+
+    @property
+    def receive_function(self) -> Optional["FunctionContract"]:
+        if self._receive_function is None:
+            for f in self.functions:
+                if f.is_receive:
+                    self._receive_function = f
+                    break
+        return self._receive_function
+
     def available_elements_from_inheritances(
         self,
         elements: Dict[str, "Function"],

slither/tools/upgradeability/checks/variables_order.py

@@ -6,6 +6,7 @@
     AbstractCheck,
     CHECK_INFO,
 )
+from slither.utils.upgradeability import get_missing_vars
 from slither.utils.output import Output
 
 
@@ -55,25 +56,13 @@ def _check(self) -> List[Output]:
         contract2 = self.contract_v2
 
         assert contract2
-
-        order1 = [
-            variable
-            for variable in contract1.state_variables_ordered
-            if not (variable.is_constant or variable.is_immutable)
-        ]
-        order2 = [
-            variable
-            for variable in contract2.state_variables_ordered
-            if not (variable.is_constant or variable.is_immutable)
-        ]
+        missing = get_missing_vars(contract1, contract2)
 
         results = []
-        for idx, _ in enumerate(order1):
-            variable1 = order1[idx]
-            if len(order2) <= idx:
-                info: CHECK_INFO = ["Variable missing in ", contract2, ": ", variable1, "\n"]
-                json = self.generate_result(info)
-                results.append(json)
+        for variable1 in missing:
+            info: CHECK_INFO = ["Variable missing in ", contract2, ": ", variable1, "\n"]
+            json = self.generate_result(info)
+            results.append(json)
 
         return results