#2586 use enum to track variable location instead of string

slither/core/variables/local_variable.py

@@ -1,11 +1,10 @@
-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
+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, VariableLocation
 
 if TYPE_CHECKING:  # type: ignore
     from slither.core.declarations import Function
@@ -14,7 +13,7 @@
 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 +24,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 +52,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:
             return False
-        if self.location == "calldata":
+        if self.location == VariableLocation.CALLDATA:
             return False
         # Use by slithIR SSA
-        if self.location == "reference_to_storage":
+        if self.location == VariableLocation.REFERENCE_TO_STORAGE:
             return False
-        if self.location == "storage":
+        if self.location == VariableLocation.STORAGE:
             return True
 
         if isinstance(self.type, (ArrayType, MappingType)):

slither/core/variables/state_variable.py

@@ -1,7 +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.variable import Variable
+from slither.core.variables.variable import Variable, VariableLocation
 
 if TYPE_CHECKING:
     from slither.core.cfg.node import Node
@@ -12,7 +12,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 +22,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 +41,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
         )
 
     @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
 
     # endregion
     ###################################################################################

slither/core/variables/variable.py

@@ -2,6 +2,7 @@
     Variable module
 """
 from typing import Optional, TYPE_CHECKING, List, Union, Tuple
+from enum import Enum
 
 from slither.core.source_mapping.source_mapping import SourceMapping
 from slither.core.solidity_types.type import Type
@@ -11,6 +12,16 @@
     from slither.core.expressions.expression import Expression
     from slither.core.declarations import Function
 
+
+class VariableLocation(Enum):
+    DEFAULT = "default"
+    MEMORY = "memory"
+    CALLDATA = "calldata"
+    STORAGE = "storage"
+    REFERENCE_TO_STORAGE = "reference_to_storage"
+    TRANSIENT = "transient"
+
+
 # pylint: disable=too-many-instance-attributes
 class Variable(SourceMapping):
     def __init__(self) -> None:

slither/detectors/compiler_bugs/array_by_reference.py

@@ -13,6 +13,7 @@
 from slither.core.solidity_types.array_type import ArrayType
 from slither.core.variables.state_variable import StateVariable
 from slither.core.variables.local_variable import LocalVariable
+from slither.core.variables.variable import VariableLocation
 from slither.core.cfg.node import Node
 from slither.core.declarations.contract import Contract
 from slither.core.declarations.function_contract import FunctionContract
@@ -84,7 +85,10 @@ def get_funcs_modifying_array_params(contracts: List[Contract]) -> Set[FunctionC
                 # Determine if this function takes an array as a parameter and the location isn't storage.
                 # If it has been written to, we know this sets an non-storage-ref array.
                 for param in function.parameters:
-                    if isinstance(param.type, ArrayType) and param.location != "storage":
+                    if (
+                        isinstance(param.type, ArrayType)
+                        and param.location != VariableLocation.STORAGE
+                    ):
                         if param in function.variables_written:
                             results.add(function)
                             break
@@ -132,8 +136,14 @@ def detect_calls_passing_ref_to_function(
                         # If it is a state variable OR a local variable referencing storage, we add it to the list.
                         if (
                             isinstance(arg, StateVariable)
-                            or (isinstance(arg, LocalVariable) and arg.location == "storage")
-                        ) and (isinstance(param.type, ArrayType) and param.location != "storage"):
+                            or (
+                                isinstance(arg, LocalVariable)
+                                and arg.location == VariableLocation.STORAGE
+                            )
+                        ) and (
+                            isinstance(param.type, ArrayType)
+                            and param.location != VariableLocation.STORAGE
+                        ):
                             results.append((ir.node, arg, ir.function))
         return results
 

slither/detectors/functions/external_function.py

@@ -4,7 +4,7 @@
 from slither.core.declarations.structure import Structure
 from slither.core.solidity_types.array_type import ArrayType
 from slither.core.solidity_types.user_defined_type import UserDefinedType
-from slither.core.variables.variable import Variable
+from slither.core.variables.variable import Variable, VariableLocation
 from slither.detectors.abstract_detector import (
     AbstractDetector,
     DetectorClassification,
@@ -170,7 +170,7 @@ def _detect(self) -> List[Output]:  # pylint: disable=too-many-locals,too-many-b
                 # If all of the function arguments are non-reference type or calldata, we skip it.
                 reference_args = []
                 for arg in function.parameters:
-                    if self.is_reference_type(arg) and arg.location == "memory":
+                    if self.is_reference_type(arg) and arg.location == VariableLocation.MEMORY:
                         reference_args.append(arg)
                 if len(reference_args) == 0:
                     continue

slither/detectors/variables/uninitialized_state_variables.py

@@ -14,6 +14,7 @@
 from slither.core.declarations.contract import Contract
 from slither.core.variables import Variable
 from slither.core.variables.state_variable import StateVariable
+from slither.core.variables.variable import VariableLocation
 from slither.detectors.abstract_detector import (
     AbstractDetector,
     DetectorClassification,
@@ -72,7 +73,7 @@ def _written_variables(contract: Contract) -> List[StateVariable]:
                         idx = 0
                         if ir.function:
                             for param in ir.function.parameters:
-                                if param.location == "storage":
+                                if param.location == VariableLocation.STORAGE:
                                     # If its a storage variable, add either the variable
                                     # Or the variable it points to if its a reference
                                     if isinstance(ir.arguments[idx], ReferenceVariable):

slither/slithir/utils/ssa.py

@@ -18,6 +18,7 @@
 from slither.core.solidity_types.type import Type
 from slither.core.variables.local_variable import LocalVariable
 from slither.core.variables.state_variable import StateVariable
+from slither.core.variables.variable import VariableLocation
 from slither.core.variables.top_level_variable import TopLevelVariable
 from slither.core.variables.variable import Variable
 from slither.slithir.exceptions import SlithIRError
@@ -166,7 +167,7 @@ def add_ssa_ir(
             if new_var.is_storage:
                 fake_variable = LocalIRVariable(v)
                 fake_variable.name = "STORAGE_" + fake_variable.name
-                fake_variable.set_location("reference_to_storage")
+                fake_variable.set_location(VariableLocation.REFERENCE_TO_STORAGE)
                 new_var.refers_to = {fake_variable}
                 init_local_variables_instances[fake_variable.name] = fake_variable
             init_local_variables_instances[v.name] = new_var
@@ -178,7 +179,7 @@ def add_ssa_ir(
             if new_var.is_storage:
                 fake_variable = LocalIRVariable(v)
                 fake_variable.name = "STORAGE_" + fake_variable.name
-                fake_variable.set_location("reference_to_storage")
+                fake_variable.set_location(VariableLocation.REFERENCE_TO_STORAGE)
                 new_var.refers_to = {fake_variable}
                 init_local_variables_instances[fake_variable.name] = fake_variable
             init_local_variables_instances[v.name] = new_var

slither/solc_parsing/declarations/custom_error.py

@@ -4,6 +4,7 @@
 from slither.core.declarations.custom_error_contract import CustomErrorContract
 from slither.core.declarations.custom_error_top_level import CustomErrorTopLevel
 from slither.core.variables.local_variable import LocalVariable
+from slither.core.variables.variable import VariableLocation
 from slither.solc_parsing.declarations.caller_context import CallerContextExpression
 from slither.solc_parsing.variables.local_variable import LocalVariableSolc
 
@@ -103,8 +104,8 @@ def _add_param(self, param: Dict) -> LocalVariableSolc:
             local_var_parser.analyze(self)
 
         # see https://solidity.readthedocs.io/en/v0.4.24/types.html?highlight=storage%20location#data-location
-        if local_var.location == "default":
-            local_var.set_location("memory")
+        if local_var.location == VariableLocation.DEFAULT:
+            local_var.set_location(VariableLocation.MEMORY)
 
         return local_var_parser
 

slither/solc_parsing/declarations/function.py

@@ -14,6 +14,7 @@
 from slither.core.source_mapping.source_mapping import Source
 from slither.core.variables.local_variable import LocalVariable
 from slither.core.variables.local_variable_init_from_tuple import LocalVariableInitFromTuple
+from slither.core.variables.variable import VariableLocation
 from slither.solc_parsing.cfg.node import NodeSolc
 from slither.solc_parsing.declarations.caller_context import CallerContextExpression
 from slither.solc_parsing.exceptions import ParsingError
@@ -1313,8 +1314,8 @@ def _add_param(self, param: Dict, initialized: bool = False) -> LocalVariableSol
             local_var.initialized = True
 
         # see https://solidity.readthedocs.io/en/v0.4.24/types.html?highlight=storage%20location#data-location
-        if local_var.location == "default":
-            local_var.set_location("memory")
+        if local_var.location == VariableLocation.DEFAULT:
+            local_var.set_location(VariableLocation.MEMORY)
 
         self._add_local_variable(local_var_parser)
         return local_var_parser

slither/solc_parsing/variables/local_variable.py

@@ -1,5 +1,6 @@
 from typing import Dict
 
+from slither.core.variables.variable import VariableLocation
 from slither.solc_parsing.variables.variable_declaration import VariableDeclarationSolc
 from slither.core.variables.local_variable import LocalVariable
 
@@ -20,14 +21,14 @@ def _analyze_variable_attributes(self, attributes: Dict) -> None:
         Can be storage/memory or default
         """
         if "storageLocation" in attributes:
-            location = attributes["storageLocation"]
+            location = VariableLocation(attributes["storageLocation"])
             self.underlying_variable.set_location(location)
         else:
             if "memory" in attributes["type"]:
-                self.underlying_variable.set_location("memory")
+                self.underlying_variable.set_location(VariableLocation.MEMORY)
             elif "storage" in attributes["type"]:
-                self.underlying_variable.set_location("storage")
+                self.underlying_variable.set_location(VariableLocation.STORAGE)
             else:
-                self.underlying_variable.set_location("default")
+                self.underlying_variable.set_location(VariableLocation.DEFAULT)
 
         super()._analyze_variable_attributes(attributes)

slither/solc_parsing/variables/local_variable_init_from_tuple.py

@@ -2,6 +2,7 @@
 
 from slither.solc_parsing.variables.variable_declaration import VariableDeclarationSolc
 from slither.core.variables.local_variable_init_from_tuple import LocalVariableInitFromTuple
+from slither.core.variables.variable import VariableLocation
 
 
 class LocalVariableInitFromTupleSolc(VariableDeclarationSolc):
@@ -27,10 +28,10 @@ def _analyze_variable_attributes(self, attributes: Dict) -> None:
             self.underlying_variable.set_location(location)
         else:
             if "memory" in attributes["type"]:
-                self.underlying_variable.set_location("memory")
+                self.underlying_variable.set_location(VariableLocation.MEMORY)
             elif "storage" in attributes["type"]:
-                self.underlying_variable.set_location("storage")
+                self.underlying_variable.set_location(VariableLocation.STORAGE)
             else:
-                self.underlying_variable.set_location("default")
+                self.underlying_variable.set_location(VariableLocation.DEFAULT)
 
         super()._analyze_variable_attributes(attributes)

slither/solc_parsing/variables/state_variable.py

@@ -2,6 +2,7 @@
 
 from slither.solc_parsing.variables.variable_declaration import VariableDeclarationSolc
 from slither.core.variables.state_variable import StateVariable
+from slither.core.variables.variable import VariableLocation
 
 
 class StateVariableSolc(VariableDeclarationSolc):
@@ -25,6 +26,6 @@ def _analyze_variable_attributes(self, attributes: Dict) -> None:
             # We don't have to support legacy ast
             # as transient location was added in 0.8.28
             # and we know it must be default
-            self.underlying_variable.set_location("default")
+            self.underlying_variable.set_location(VariableLocation.DEFAULT)
 
         super()._analyze_variable_attributes(attributes)

slither/solc_parsing/yul/parse_yul.py

@@ -28,6 +28,7 @@
 from slither.core.solidity_types import ElementaryType
 from slither.core.source_mapping.source_mapping import SourceMapping
 from slither.core.variables.local_variable import LocalVariable
+from slither.core.variables.variable import VariableLocation
 from slither.exceptions import SlitherException
 from slither.solc_parsing.yul.evm_functions import (
     format_function_descriptor,
@@ -203,7 +204,7 @@ def __init__(self, var: LocalVariable, root: YulScope, ast: Dict) -> None:
 
         var.name = _name_to_yul_name(ast["name"], root.id)
         var.set_type(ElementaryType("uint256"))
-        var.set_location("memory")
+        var.set_location(VariableLocation.MEMORY)
 
     @property
     def underlying(self) -> LocalVariable:
@@ -778,15 +779,15 @@ def _parse_yul_magic_suffixes(name: str, root: YulScope) -> Optional[Expression]
         if variable_found:
             return variable_found
         var = root.function.get_local_variable_from_name(potential_name)
-        if var and var.location == "calldata":
+        if var and var.location == VariableLocation.CALLDATA:
             return Identifier(var)
     if name.endswith(".length"):
         # TODO: length should create a new IP operation LENGTH var
         # The code below is an hotfix to allow slither to process length in yul
         # Until we have a better support
         potential_name = name[:-7]
         var = root.function.get_local_variable_from_name(potential_name)
-        if var and var.location == "calldata":
+        if var and var.location == VariableLocation.CALLDATA:
             return Identifier(var)
     return None