#2586 use enum to track variable location instead of string

slither/core/variables/local_variable.py

@@ -1,20 +1,28 @@
-from typing import Optional, TYPE_CHECKING
-
-from slither.core.variables.variable import Variable
-from slither.core.solidity_types.user_defined_type import UserDefinedType
-from slither.core.solidity_types.mapping_type import MappingType
-from slither.core.solidity_types.elementary_type import ElementaryType
+import enum
+from typing import TYPE_CHECKING, Optional
 
 from slither.core.declarations.structure import Structure
+from slither.core.solidity_types.elementary_type import ElementaryType
+from slither.core.solidity_types.mapping_type import MappingType
+from slither.core.solidity_types.user_defined_type import UserDefinedType
+from slither.core.variables.variable import Variable
 
 if TYPE_CHECKING:  # type: ignore
     from slither.core.declarations import Function
 
 
+class VariableLocation(enum.Enum):
+    MEMORY = "memory"
+    CALLDATA = "calldata"
+    STORAGE = "storage"
+    REFERENCE_TO_STORAGE = "reference_to_storage"
+    TRANSIENT = "transient"
+
+
 class LocalVariable(Variable):
     def __init__(self) -> None:
         super().__init__()
-        self._location: Optional[str] = None
+        self._location: Optional[VariableLocation] = None
         self._function: Optional["Function"] = None
 
     def set_function(self, function: "Function") -> None:
@@ -25,16 +33,16 @@ def function(self) -> "Function":
         assert self._function
         return self._function
 
-    def set_location(self, loc: str) -> None:
+    def set_location(self, loc: VariableLocation) -> None:
         self._location = loc
 
     @property
-    def location(self) -> Optional[str]:
+    def location(self) -> Optional[VariableLocation]:
         """
             Variable Location
             Can be storage/memory or default
         Returns:
-            (str)
+            (VariableLocation)
         """
         return self._location
 
@@ -53,14 +61,14 @@ def is_storage(self) -> bool:
         # pylint: disable=import-outside-toplevel
         from slither.core.solidity_types.array_type import ArrayType
 
-        if self.location == "memory":
+        if self.location == VariableLocation.MEMORY.value:
             return False
-        if self.location == "calldata":
+        if self.location == VariableLocation.CALLDATA.value:
             return False
         # Use by slithIR SSA
-        if self.location == "reference_to_storage":
+        if self.location == VariableLocation.REFERENCE_TO_STORAGE.value:
             return False
-        if self.location == "storage":
+        if self.location == VariableLocation.STORAGE.value:
             return True
 
         if isinstance(self.type, (ArrayType, MappingType)):

slither/core/variables/state_variable.py

@@ -1,6 +1,7 @@
-from typing import Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 from slither.core.declarations.contract_level import ContractLevel
+from slither.core.variables.local_variable import VariableLocation
 from slither.core.variables.variable import Variable
 
 if TYPE_CHECKING:
@@ -12,7 +13,7 @@ class StateVariable(ContractLevel, Variable):
     def __init__(self) -> None:
         super().__init__()
         self._node_initialization: Optional["Node"] = None
-        self._location: Optional[str] = None
+        self._location: Optional[VariableLocation] = None
 
     def is_declared_by(self, contract: "Contract") -> bool:
         """
@@ -22,16 +23,16 @@ def is_declared_by(self, contract: "Contract") -> bool:
         """
         return self.contract == contract
 
-    def set_location(self, loc: str) -> None:
+    def set_location(self, loc: VariableLocation) -> None:
         self._location = loc
 
     @property
-    def location(self) -> Optional[str]:
+    def location(self) -> Optional[VariableLocation]:
         """
             Variable Location
             Can be default or transient
         Returns:
-            (str)
+            (VariableLocation)
         """
         return self._location
 
@@ -41,15 +42,17 @@ def is_stored(self) -> bool:
         Checks if the state variable is stored, based on it not being constant or immutable or transient.
         """
         return (
-            not self._is_constant and not self._is_immutable and not self._location == "transient"
+            not self._is_constant
+            and not self._is_immutable
+            and not self._location == VariableLocation.TRANSIENT.value
         )
 
     @property
     def is_transient(self) -> bool:
         """
         Checks if the state variable is transient. A transient variable can not be constant or immutable.
         """
-        return self._location == "transient"
+        return self._location == VariableLocation.TRANSIENT.value
 
     # endregion
     ###################################################################################

tests/unit/slithir/test_ssa_generation.py

@@ -3,36 +3,36 @@
 from argparse import ArgumentTypeError
 from collections import defaultdict
 from inspect import getsourcefile
-from typing import Union, List, Dict, Callable
+from typing import Callable, Dict, List, Union
 
 import pytest
 from solc_select.solc_select import valid_version as solc_valid_version
 
 from slither import Slither
 from slither.core.cfg.node import Node, NodeType
-from slither.core.declarations import Function, Contract
+from slither.core.declarations import Contract, Function
 from slither.core.solidity_types import ArrayType, ElementaryType
-from slither.core.variables.local_variable import LocalVariable
+from slither.core.variables.local_variable import LocalVariable, VariableLocation
 from slither.core.variables.state_variable import StateVariable
 from slither.slithir.operations import (
-    OperationWithLValue,
-    Phi,
     Assignment,
-    HighLevelCall,
-    Return,
-    Operation,
     Binary,
     BinaryType,
-    InternalCall,
+    HighLevelCall,
     Index,
     InitArray,
+    InternalCall,
     NewArray,
+    Operation,
+    OperationWithLValue,
+    Phi,
+    Return,
 )
 from slither.slithir.utils.ssa import is_used_later
 from slither.slithir.variables import (
     Constant,
-    ReferenceVariable,
     LocalIRVariable,
+    ReferenceVariable,
     StateIRVariable,
     TemporaryVariableSSA,
 )
@@ -120,7 +120,7 @@ def ssa_basic_properties(function: Function) -> None:
         if v and v.name:
             ssa_defs[v.name] += 1
 
-    for (k, count) in lvalue_assignments.items():
+    for k, count in lvalue_assignments.items():
         assert ssa_defs[k] >= count
 
     # Helper 5/6
@@ -134,7 +134,7 @@ def check_property_5_and_6(
             assert var.is_storage == ssa_var.is_storage
             if ssa_var.is_storage:
                 assert len(ssa_var.refers_to) == 1
-                assert ssa_var.refers_to[0].location == "reference_to_storage"
+                assert ssa_var.refers_to[0].location == VariableLocation.REFERENCE_TO_STORAGE.value
 
     # 5
     check_property_5_and_6(function.parameters, function.parameters_ssa)