slither.utils.code_generation

  1# Functions for generating Solidity code
  2from typing import TYPE_CHECKING, Optional
  3
  4from slither.utils.type import (
  5    convert_type_for_solidity_signature_to_string,
  6    export_nested_types_from_variable,
  7    export_return_type_from_variable,
  8)
  9from slither.core.solidity_types import (
 10    Type,
 11    UserDefinedType,
 12    MappingType,
 13    ArrayType,
 14    ElementaryType,
 15    TypeAlias,
 16)
 17from slither.core.declarations import Structure, StructureContract, Enum, Contract
 18
 19if TYPE_CHECKING:
 20    from slither.core.declarations import FunctionContract, CustomErrorContract
 21    from slither.core.variables.state_variable import StateVariable
 22    from slither.core.variables.local_variable import LocalVariable
 23    from slither.core.variables.structure_variable import StructureVariable
 24
 25
 26# pylint: disable=too-many-arguments,too-many-locals,too-many-branches
 27def generate_interface(
 28    contract: "Contract",
 29    unroll_structs: bool = True,
 30    include_events: bool = True,
 31    include_errors: bool = True,
 32    include_enums: bool = True,
 33    include_structs: bool = True,
 34) -> str:
 35    """
 36    Generates code for a Solidity interface to the contract.
 37    Args:
 38        contract: A Contract object.
 39        unroll_structs: Whether to use structures' underlying types instead of the user-defined type (default: True).
 40        include_events: Whether to include event signatures in the interface (default: True).
 41        include_errors: Whether to include custom error signatures in the interface (default: True).
 42        include_enums: Whether to include enum definitions in the interface (default: True).
 43        include_structs: Whether to include struct definitions in the interface (default: True).
 44
 45    Returns:
 46        A string with the code for an interface, with function stubs for all public or external functions and
 47        state variables, as well as any events, custom errors and/or structs declared in the contract.
 48    """
 49    interface = f"interface I{contract.name} {{\n"
 50    if include_events:
 51        for event in contract.events:
 52            name, args = event.signature
 53            interface += f"    event {name}({', '.join(args)});\n"
 54    if include_errors:
 55        for error in contract.custom_errors:
 56            interface += f"    error {generate_custom_error_interface(error, unroll_structs)};\n"
 57    if include_enums:
 58        for enum in contract.enums:
 59            interface += f"    enum {enum.name} {{ {', '.join(enum.values)} }}\n"
 60    if include_structs:
 61        # Include structures defined in this contract and at the top level
 62        structs = contract.structures + contract.compilation_unit.structures_top_level
 63        # Function signatures may reference other structures as well
 64        # Include structures defined in libraries used for them
 65        for _for in contract.using_for.keys():
 66            if (
 67                isinstance(_for, UserDefinedType)
 68                and isinstance(_for.type, StructureContract)
 69                and _for.type not in structs
 70            ):
 71                structs.append(_for.type)
 72        # Include any other structures used as function arguments/returns
 73        for func in contract.functions_entry_points:
 74            for arg in func.parameters + func.returns:
 75                _type = arg.type
 76                if isinstance(_type, ArrayType):
 77                    _type = _type.type
 78                while isinstance(_type, MappingType):
 79                    _type = _type.type_to
 80                if isinstance(_type, UserDefinedType):
 81                    _type = _type.type
 82                if isinstance(_type, Structure) and _type not in structs:
 83                    structs.append(_type)
 84        for struct in structs:
 85            interface += generate_struct_interface_str(struct, indent=4)
 86            for elem in struct.elems_ordered:
 87                if (
 88                    isinstance(elem.type, UserDefinedType)
 89                    and isinstance(elem.type.type, StructureContract)
 90                    and elem.type.type not in structs
 91                ):
 92                    structs.append(elem.type.type)
 93    for var in contract.state_variables_entry_points:
 94        # if any(func.name == var.name for func in contract.functions_entry_points):
 95        #     # ignore public variables that override a public function
 96        #     continue
 97        var_sig = generate_interface_variable_signature(var, unroll_structs)
 98        if var_sig is not None and var_sig != "":
 99            interface += f"    function {var_sig};\n"
100    for func in contract.functions_entry_points:
101        if func.is_constructor or func.is_fallback or func.is_receive or not func.is_implemented:
102            continue
103        interface += (
104            f"    function {generate_interface_function_signature(func, unroll_structs)};\n"
105        )
106    interface += "}\n\n"
107    return interface
108
109
110def generate_interface_variable_signature(
111    var: "StateVariable", unroll_structs: bool = True
112) -> Optional[str]:
113    if var.visibility in ["private", "internal"]:
114        return None
115    if isinstance(var.type, UserDefinedType) and isinstance(var.type.type, Structure):
116        for elem in var.type.type.elems_ordered:
117            if isinstance(elem.type, MappingType):
118                return ""
119    if unroll_structs:
120        params = [
121            convert_type_for_solidity_signature_to_string(x).replace("(", "").replace(")", "")
122            for x in export_nested_types_from_variable(var)
123        ]
124        returns = [
125            convert_type_for_solidity_signature_to_string(x).replace("(", "").replace(")", "")
126            for x in export_return_type_from_variable(var)
127        ]
128    else:
129        _, params, _ = var.signature
130        params = [p + " memory" if p in ["bytes", "string"] else p for p in params]
131        returns = []
132        _type = var.type
133        while isinstance(_type, MappingType):
134            _type = _type.type_to
135        while isinstance(_type, (ArrayType, UserDefinedType)):
136            _type = _type.type
137        if isinstance(_type, TypeAlias):
138            _type = _type.type
139        if isinstance(_type, Structure):
140            if any(isinstance(elem.type, MappingType) for elem in _type.elems_ordered):
141                return ""
142        ret = str(_type)
143        if isinstance(_type, Structure) or (isinstance(_type, Type) and _type.is_dynamic):
144            ret += " memory"
145        elif isinstance(_type, Contract):
146            ret = "address"
147        returns.append(ret)
148    return f"{var.name}({','.join(params)}) external returns ({', '.join(returns)})"
149
150
151def generate_interface_function_signature(
152    func: "FunctionContract", unroll_structs: bool = True
153) -> Optional[str]:
154    """
155    Generates a string of the form:
156        func_name(type1,type2) external {payable/view/pure} returns (type3)
157
158    Args:
159        func: A FunctionContract object
160        unroll_structs: Determines whether structs are unrolled into underlying types (default: True)
161
162    Returns:
163        The function interface as a str (contains the return values).
164        Returns None if the function is private or internal, or is a constructor/fallback/receive.
165    """
166
167    def format_var(var: "LocalVariable", unroll: bool) -> str:
168        if unroll:
169            return (
170                convert_type_for_solidity_signature_to_string(var.type)
171                .replace("(", "")
172                .replace(")", "")
173            )
174        if var.type.is_dynamic:
175            return f"{_handle_dynamic_struct_elem(var.type)} {var.location}"
176        if isinstance(var.type, ArrayType) and isinstance(
177            var.type.type, (UserDefinedType, ElementaryType)
178        ):
179            return (
180                convert_type_for_solidity_signature_to_string(var.type)
181                .replace("(", "")
182                .replace(")", "")
183                + f" {var.location}"
184            )
185        if isinstance(var.type, UserDefinedType):
186            if isinstance(var.type.type, Structure):
187                return f"{str(var.type.type)} memory"
188            if isinstance(var.type.type, Enum):
189                return str(var.type.type)
190            if isinstance(var.type.type, Contract):
191                return "address"
192        if isinstance(var.type, TypeAlias):
193            return str(var.type.type)
194        return str(var.type)
195
196    name, _, _ = func.signature
197    if (
198        func not in func.contract.functions_entry_points
199        or func.is_constructor
200        or func.is_fallback
201        or func.is_receive
202    ):
203        return None
204    view = " view" if func.view and not func.pure else ""
205    pure = " pure" if func.pure else ""
206    payable = " payable" if func.payable else ""
207    # Make sure the function doesn't return a struct with nested mappings
208    for ret in func.returns:
209        if isinstance(ret.type, UserDefinedType) and isinstance(ret.type.type, Structure):
210            for elem in ret.type.type.elems_ordered:
211                if isinstance(elem.type, MappingType):
212                    return ""
213    returns = [format_var(ret, unroll_structs) for ret in func.returns]
214    parameters = [format_var(param, unroll_structs) for param in func.parameters]
215    _interface_signature_str = (
216        name + "(" + ",".join(parameters) + ") external" + payable + pure + view
217    )
218    if len(returns) > 0:
219        _interface_signature_str += " returns (" + ",".join(returns) + ")"
220    return _interface_signature_str
221
222
223def generate_struct_interface_str(struct: "Structure", indent: int = 0) -> str:
224    """
225    Generates code for a structure declaration in an interface of the form:
226        struct struct_name {
227            elem1_type elem1_name;
228            elem2_type elem2_name;
229            ...        ...
230        }
231    Args:
232        struct: A Structure object.
233        indent: Number of spaces to indent the code block with.
234
235    Returns:
236        The structure declaration code as a string.
237    """
238    spaces = ""
239    for _ in range(0, indent):
240        spaces += " "
241    definition = f"{spaces}struct {struct.name} {{\n"
242    for elem in struct.elems_ordered:
243        if elem.type.is_dynamic:
244            definition += f"{spaces}    {_handle_dynamic_struct_elem(elem.type)} {elem.name};\n"
245        elif isinstance(elem.type, UserDefinedType):
246            if isinstance(elem.type.type, Structure):
247                definition += f"{spaces}    {elem.type.type} {elem.name};\n"
248            else:
249                definition += f"{spaces}    {convert_type_for_solidity_signature_to_string(elem.type)} {elem.name};\n"
250        elif isinstance(elem.type, TypeAlias):
251            definition += f"{spaces}    {elem.type.type} {elem.name};\n"
252        else:
253            definition += f"{spaces}    {elem.type} {elem.name};\n"
254    definition += f"{spaces}}}\n"
255    return definition
256
257
258def _handle_dynamic_struct_elem(elem_type: Type) -> str:
259    assert elem_type.is_dynamic
260    if isinstance(elem_type, ElementaryType):
261        return f"{elem_type}"
262    if isinstance(elem_type, ArrayType):
263        base_type = elem_type.type
264        if isinstance(base_type, UserDefinedType):
265            if isinstance(base_type.type, Contract):
266                return "address[]"
267            if isinstance(base_type.type, Enum):
268                return convert_type_for_solidity_signature_to_string(elem_type)
269            return f"{base_type.type.name}[]"
270        return f"{base_type}[]"
271    if isinstance(elem_type, MappingType):
272        type_to = elem_type.type_to
273        type_from = elem_type.type_from
274        if isinstance(type_from, UserDefinedType) and isinstance(type_from.type, Contract):
275            type_from = ElementaryType("address")
276        if isinstance(type_to, MappingType):
277            return f"mapping({type_from} => {_handle_dynamic_struct_elem(type_to)})"
278        if isinstance(type_to, UserDefinedType):
279            if isinstance(type_to.type, Contract):
280                return f"mapping({type_from} => address)"
281            return f"mapping({type_from} => {type_to.type.name})"
282        return f"{elem_type}"
283    return ""
284
285
286def generate_custom_error_interface(
287    error: "CustomErrorContract", unroll_structs: bool = True
288) -> str:
289    args = [
290        convert_type_for_solidity_signature_to_string(arg.type).replace("(", "").replace(")", "")
291        if unroll_structs
292        else str(arg.type.type)
293        if isinstance(arg.type, UserDefinedType) and isinstance(arg.type.type, (Structure, Enum))
294        else str(arg.type)
295        for arg in error.parameters
296    ]
297    return f"{error.name}({', '.join(args)})"