647 lines
17 KiB
Python
647 lines
17 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
import copy
|
|
import json
|
|
import re
|
|
import subprocess
|
|
from enum import Enum as PyEnum
|
|
from pathlib import Path
|
|
from typing import Callable
|
|
from urllib import request
|
|
|
|
VoidFn = Callable[[], None]
|
|
|
|
CHEATCODES_JSON_URL = "https://raw.githubusercontent.com/foundry-rs/foundry/master/crates/cheatcodes/assets/cheatcodes.json"
|
|
OUT_PATH = "src/Vm.sol"
|
|
|
|
VM_SAFE_DOC = """\
|
|
/// The `VmSafe` interface does not allow manipulation of the EVM state or other actions that may
|
|
/// result in Script simulations differing from on-chain execution. It is recommended to only use
|
|
/// these cheats in scripts.
|
|
"""
|
|
|
|
VM_DOC = """\
|
|
/// The `Vm` interface does allow manipulation of the EVM state. These are all intended to be used
|
|
/// in tests, but it is not recommended to use these cheats in scripts.
|
|
"""
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Generate Vm.sol based on the cheatcodes json created by Foundry")
|
|
parser.add_argument(
|
|
"--from",
|
|
metavar="PATH",
|
|
dest="path",
|
|
required=False,
|
|
help="path to a json file containing the Vm interface, as generated by Foundry")
|
|
args = parser.parse_args()
|
|
json_str = request.urlopen(CHEATCODES_JSON_URL).read().decode("utf-8") if args.path is None else Path(args.path).read_text()
|
|
contract = Cheatcodes.from_json(json_str)
|
|
|
|
ccs = contract.cheatcodes
|
|
ccs = list(filter(lambda cc: cc.status not in ["experimental", "internal"], ccs))
|
|
ccs.sort(key=lambda cc: cc.func.id)
|
|
|
|
safe = list(filter(lambda cc: cc.safety == "safe", ccs))
|
|
safe.sort(key=CmpCheatcode)
|
|
unsafe = list(filter(lambda cc: cc.safety == "unsafe", ccs))
|
|
unsafe.sort(key=CmpCheatcode)
|
|
assert len(safe) + len(unsafe) == len(ccs)
|
|
|
|
prefix_with_group_headers(safe)
|
|
prefix_with_group_headers(unsafe)
|
|
|
|
out = ""
|
|
|
|
out += "// Automatically @generated by scripts/vm.py. Do not modify manually.\n\n"
|
|
|
|
pp = CheatcodesPrinter(
|
|
spdx_identifier="MIT OR Apache-2.0",
|
|
solidity_requirement=">=0.6.2 <0.9.0",
|
|
abicoder_pragma=True,
|
|
)
|
|
pp.p_prelude()
|
|
pp.prelude = False
|
|
out += pp.finish()
|
|
|
|
out += "\n\n"
|
|
out += VM_SAFE_DOC
|
|
vm_safe = Cheatcodes(
|
|
# TODO: Custom errors were introduced in 0.8.4
|
|
errors=[], # contract.errors
|
|
events=contract.events,
|
|
enums=contract.enums,
|
|
structs=contract.structs,
|
|
cheatcodes=safe,
|
|
)
|
|
pp.p_contract(vm_safe, "VmSafe")
|
|
out += pp.finish()
|
|
|
|
out += "\n\n"
|
|
out += VM_DOC
|
|
vm_unsafe = Cheatcodes(
|
|
errors=[],
|
|
events=[],
|
|
enums=[],
|
|
structs=[],
|
|
cheatcodes=unsafe,
|
|
)
|
|
pp.p_contract(vm_unsafe, "Vm", "VmSafe")
|
|
out += pp.finish()
|
|
|
|
# Compatibility with <0.8.0
|
|
def memory_to_calldata(m: re.Match) -> str:
|
|
return " calldata " + m.group(1)
|
|
|
|
out = re.sub(r" memory (.*returns)", memory_to_calldata, out)
|
|
|
|
with open(OUT_PATH, "w") as f:
|
|
f.write(out)
|
|
|
|
forge_fmt = ["forge", "fmt", OUT_PATH]
|
|
res = subprocess.run(forge_fmt)
|
|
assert res.returncode == 0, f"command failed: {forge_fmt}"
|
|
|
|
print(f"Wrote to {OUT_PATH}")
|
|
|
|
|
|
class CmpCheatcode:
|
|
cheatcode: "Cheatcode"
|
|
|
|
def __init__(self, cheatcode: "Cheatcode"):
|
|
self.cheatcode = cheatcode
|
|
|
|
def __lt__(self, other: "CmpCheatcode") -> bool:
|
|
return cmp_cheatcode(self.cheatcode, other.cheatcode) < 0
|
|
|
|
def __eq__(self, other: "CmpCheatcode") -> bool:
|
|
return cmp_cheatcode(self.cheatcode, other.cheatcode) == 0
|
|
|
|
def __gt__(self, other: "CmpCheatcode") -> bool:
|
|
return cmp_cheatcode(self.cheatcode, other.cheatcode) > 0
|
|
|
|
|
|
def cmp_cheatcode(a: "Cheatcode", b: "Cheatcode") -> int:
|
|
if a.group != b.group:
|
|
return -1 if a.group < b.group else 1
|
|
if a.status != b.status:
|
|
return -1 if a.status < b.status else 1
|
|
if a.safety != b.safety:
|
|
return -1 if a.safety < b.safety else 1
|
|
if a.func.id != b.func.id:
|
|
return -1 if a.func.id < b.func.id else 1
|
|
return 0
|
|
|
|
|
|
# HACK: A way to add group header comments without having to modify printer code
|
|
def prefix_with_group_headers(cheats: list["Cheatcode"]):
|
|
s = set()
|
|
for i, cheat in enumerate(cheats):
|
|
if cheat.group in s:
|
|
continue
|
|
|
|
s.add(cheat.group)
|
|
|
|
c = copy.deepcopy(cheat)
|
|
c.func.description = ""
|
|
c.func.declaration = f"// ======== {group(c.group)} ========"
|
|
cheats.insert(i, c)
|
|
return cheats
|
|
|
|
|
|
def group(s: str) -> str:
|
|
if s == "evm":
|
|
return "EVM"
|
|
if s == "json":
|
|
return "JSON"
|
|
return s[0].upper() + s[1:]
|
|
|
|
|
|
class Visibility(PyEnum):
|
|
EXTERNAL: str = "external"
|
|
PUBLIC: str = "public"
|
|
INTERNAL: str = "internal"
|
|
PRIVATE: str = "private"
|
|
|
|
def __str__(self):
|
|
return self.value
|
|
|
|
|
|
class Mutability(PyEnum):
|
|
PURE: str = "pure"
|
|
VIEW: str = "view"
|
|
NONE: str = ""
|
|
|
|
def __str__(self):
|
|
return self.value
|
|
|
|
|
|
class Function:
|
|
id: str
|
|
description: str
|
|
declaration: str
|
|
visibility: Visibility
|
|
mutability: Mutability
|
|
signature: str
|
|
selector: str
|
|
selector_bytes: bytes
|
|
|
|
def __init__(
|
|
self,
|
|
id: str,
|
|
description: str,
|
|
declaration: str,
|
|
visibility: Visibility,
|
|
mutability: Mutability,
|
|
signature: str,
|
|
selector: str,
|
|
selector_bytes: bytes,
|
|
):
|
|
self.id = id
|
|
self.description = description
|
|
self.declaration = declaration
|
|
self.visibility = visibility
|
|
self.mutability = mutability
|
|
self.signature = signature
|
|
self.selector = selector
|
|
self.selector_bytes = selector_bytes
|
|
|
|
@staticmethod
|
|
def from_dict(d: dict) -> "Function":
|
|
return Function(
|
|
d["id"],
|
|
d["description"],
|
|
d["declaration"],
|
|
Visibility(d["visibility"]),
|
|
Mutability(d["mutability"]),
|
|
d["signature"],
|
|
d["selector"],
|
|
bytes(d["selectorBytes"]),
|
|
)
|
|
|
|
|
|
class Cheatcode:
|
|
func: Function
|
|
group: str
|
|
status: str
|
|
safety: str
|
|
|
|
def __init__(self, func: Function, group: str, status: str, safety: str):
|
|
self.func = func
|
|
self.group = group
|
|
self.status = status
|
|
self.safety = safety
|
|
|
|
@staticmethod
|
|
def from_dict(d: dict) -> "Cheatcode":
|
|
return Cheatcode(
|
|
Function.from_dict(d["func"]),
|
|
str(d["group"]),
|
|
str(d["status"]),
|
|
str(d["safety"]),
|
|
)
|
|
|
|
|
|
class Error:
|
|
name: str
|
|
description: str
|
|
declaration: str
|
|
|
|
def __init__(self, name: str, description: str, declaration: str):
|
|
self.name = name
|
|
self.description = description
|
|
self.declaration = declaration
|
|
|
|
@staticmethod
|
|
def from_dict(d: dict) -> "Error":
|
|
return Error(**d)
|
|
|
|
|
|
class Event:
|
|
name: str
|
|
description: str
|
|
declaration: str
|
|
|
|
def __init__(self, name: str, description: str, declaration: str):
|
|
self.name = name
|
|
self.description = description
|
|
self.declaration = declaration
|
|
|
|
@staticmethod
|
|
def from_dict(d: dict) -> "Event":
|
|
return Event(**d)
|
|
|
|
|
|
class EnumVariant:
|
|
name: str
|
|
description: str
|
|
|
|
def __init__(self, name: str, description: str):
|
|
self.name = name
|
|
self.description = description
|
|
|
|
|
|
class Enum:
|
|
name: str
|
|
description: str
|
|
variants: list[EnumVariant]
|
|
|
|
def __init__(self, name: str, description: str, variants: list[EnumVariant]):
|
|
self.name = name
|
|
self.description = description
|
|
self.variants = variants
|
|
|
|
@staticmethod
|
|
def from_dict(d: dict) -> "Enum":
|
|
return Enum(
|
|
d["name"],
|
|
d["description"],
|
|
list(map(lambda v: EnumVariant(**v), d["variants"])),
|
|
)
|
|
|
|
|
|
class StructField:
|
|
name: str
|
|
ty: str
|
|
description: str
|
|
|
|
def __init__(self, name: str, ty: str, description: str):
|
|
self.name = name
|
|
self.ty = ty
|
|
self.description = description
|
|
|
|
|
|
class Struct:
|
|
name: str
|
|
description: str
|
|
fields: list[StructField]
|
|
|
|
def __init__(self, name: str, description: str, fields: list[StructField]):
|
|
self.name = name
|
|
self.description = description
|
|
self.fields = fields
|
|
|
|
@staticmethod
|
|
def from_dict(d: dict) -> "Struct":
|
|
return Struct(
|
|
d["name"],
|
|
d["description"],
|
|
list(map(lambda f: StructField(**f), d["fields"])),
|
|
)
|
|
|
|
|
|
class Cheatcodes:
|
|
errors: list[Error]
|
|
events: list[Event]
|
|
enums: list[Enum]
|
|
structs: list[Struct]
|
|
cheatcodes: list[Cheatcode]
|
|
|
|
def __init__(
|
|
self,
|
|
errors: list[Error],
|
|
events: list[Event],
|
|
enums: list[Enum],
|
|
structs: list[Struct],
|
|
cheatcodes: list[Cheatcode],
|
|
):
|
|
self.errors = errors
|
|
self.events = events
|
|
self.enums = enums
|
|
self.structs = structs
|
|
self.cheatcodes = cheatcodes
|
|
|
|
@staticmethod
|
|
def from_dict(d: dict) -> "Cheatcodes":
|
|
return Cheatcodes(
|
|
errors=[Error.from_dict(e) for e in d["errors"]],
|
|
events=[Event.from_dict(e) for e in d["events"]],
|
|
enums=[Enum.from_dict(e) for e in d["enums"]],
|
|
structs=[Struct.from_dict(e) for e in d["structs"]],
|
|
cheatcodes=[Cheatcode.from_dict(e) for e in d["cheatcodes"]],
|
|
)
|
|
|
|
@staticmethod
|
|
def from_json(s) -> "Cheatcodes":
|
|
return Cheatcodes.from_dict(json.loads(s))
|
|
|
|
@staticmethod
|
|
def from_json_file(file_path: str) -> "Cheatcodes":
|
|
with open(file_path, "r") as f:
|
|
return Cheatcodes.from_dict(json.load(f))
|
|
|
|
|
|
class Item(PyEnum):
|
|
ERROR: str = "error"
|
|
EVENT: str = "event"
|
|
ENUM: str = "enum"
|
|
STRUCT: str = "struct"
|
|
FUNCTION: str = "function"
|
|
|
|
|
|
class ItemOrder:
|
|
_list: list[Item]
|
|
|
|
def __init__(self, list: list[Item]) -> None:
|
|
assert len(list) <= len(Item), "list must not contain more items than Item"
|
|
assert len(list) == len(set(list)), "list must not contain duplicates"
|
|
self._list = list
|
|
pass
|
|
|
|
def get_list(self) -> list[Item]:
|
|
return self._list
|
|
|
|
@staticmethod
|
|
def default() -> "ItemOrder":
|
|
return ItemOrder(
|
|
[
|
|
Item.ERROR,
|
|
Item.EVENT,
|
|
Item.ENUM,
|
|
Item.STRUCT,
|
|
Item.FUNCTION,
|
|
]
|
|
)
|
|
|
|
|
|
class CheatcodesPrinter:
|
|
buffer: str
|
|
|
|
prelude: bool
|
|
spdx_identifier: str
|
|
solidity_requirement: str
|
|
abicoder_v2: bool
|
|
|
|
block_doc_style: bool
|
|
|
|
indent_level: int
|
|
_indent_str: str
|
|
|
|
nl_str: str
|
|
|
|
items_order: ItemOrder
|
|
|
|
def __init__(
|
|
self,
|
|
buffer: str = "",
|
|
prelude: bool = True,
|
|
spdx_identifier: str = "UNLICENSED",
|
|
solidity_requirement: str = "",
|
|
abicoder_pragma: bool = False,
|
|
block_doc_style: bool = False,
|
|
indent_level: int = 0,
|
|
indent_with: int | str = 4,
|
|
nl_str: str = "\n",
|
|
items_order: ItemOrder = ItemOrder.default(),
|
|
):
|
|
self.prelude = prelude
|
|
self.spdx_identifier = spdx_identifier
|
|
self.solidity_requirement = solidity_requirement
|
|
self.abicoder_v2 = abicoder_pragma
|
|
self.block_doc_style = block_doc_style
|
|
self.buffer = buffer
|
|
self.indent_level = indent_level
|
|
self.nl_str = nl_str
|
|
|
|
if isinstance(indent_with, int):
|
|
assert indent_with >= 0
|
|
self._indent_str = " " * indent_with
|
|
elif isinstance(indent_with, str):
|
|
self._indent_str = indent_with
|
|
else:
|
|
assert False, "indent_with must be int or str"
|
|
|
|
self.items_order = items_order
|
|
|
|
def finish(self) -> str:
|
|
ret = self.buffer.rstrip()
|
|
self.buffer = ""
|
|
return ret
|
|
|
|
def p_contract(self, contract: Cheatcodes, name: str, inherits: str = ""):
|
|
if self.prelude:
|
|
self.p_prelude(contract)
|
|
|
|
self._p_str("interface ")
|
|
name = name.strip()
|
|
if name != "":
|
|
self._p_str(name)
|
|
self._p_str(" ")
|
|
if inherits != "":
|
|
self._p_str("is ")
|
|
self._p_str(inherits)
|
|
self._p_str(" ")
|
|
self._p_str("{")
|
|
self._p_nl()
|
|
self._with_indent(lambda: self._p_items(contract))
|
|
self._p_str("}")
|
|
self._p_nl()
|
|
|
|
def _p_items(self, contract: Cheatcodes):
|
|
for item in self.items_order.get_list():
|
|
if item == Item.ERROR:
|
|
self.p_errors(contract.errors)
|
|
elif item == Item.EVENT:
|
|
self.p_events(contract.events)
|
|
elif item == Item.ENUM:
|
|
self.p_enums(contract.enums)
|
|
elif item == Item.STRUCT:
|
|
self.p_structs(contract.structs)
|
|
elif item == Item.FUNCTION:
|
|
self.p_functions(contract.cheatcodes)
|
|
else:
|
|
assert False, f"unknown item {item}"
|
|
|
|
def p_prelude(self, contract: Cheatcodes | None = None):
|
|
self._p_str(f"// SPDX-License-Identifier: {self.spdx_identifier}")
|
|
self._p_nl()
|
|
|
|
if self.solidity_requirement != "":
|
|
req = self.solidity_requirement
|
|
elif contract and len(contract.errors) > 0:
|
|
req = ">=0.8.4 <0.9.0"
|
|
else:
|
|
req = ">=0.6.0 <0.9.0"
|
|
self._p_str(f"pragma solidity {req};")
|
|
self._p_nl()
|
|
|
|
if self.abicoder_v2:
|
|
self._p_str("pragma experimental ABIEncoderV2;")
|
|
self._p_nl()
|
|
|
|
self._p_nl()
|
|
|
|
def p_errors(self, errors: list[Error]):
|
|
for error in errors:
|
|
self._p_line(lambda: self.p_error(error))
|
|
|
|
def p_error(self, error: Error):
|
|
self._p_comment(error.description, doc=True)
|
|
self._p_line(lambda: self._p_str(error.declaration))
|
|
|
|
def p_events(self, events: list[Event]):
|
|
for event in events:
|
|
self._p_line(lambda: self.p_event(event))
|
|
|
|
def p_event(self, event: Event):
|
|
self._p_comment(event.description, doc=True)
|
|
self._p_line(lambda: self._p_str(event.declaration))
|
|
|
|
def p_enums(self, enums: list[Enum]):
|
|
for enum in enums:
|
|
self._p_line(lambda: self.p_enum(enum))
|
|
|
|
def p_enum(self, enum: Enum):
|
|
self._p_comment(enum.description, doc=True)
|
|
self._p_line(lambda: self._p_str(f"enum {enum.name} {{"))
|
|
self._with_indent(lambda: self.p_enum_variants(enum.variants))
|
|
self._p_line(lambda: self._p_str("}"))
|
|
|
|
def p_enum_variants(self, variants: list[EnumVariant]):
|
|
for i, variant in enumerate(variants):
|
|
self._p_indent()
|
|
self._p_comment(variant.description)
|
|
|
|
self._p_indent()
|
|
self._p_str(variant.name)
|
|
if i < len(variants) - 1:
|
|
self._p_str(",")
|
|
self._p_nl()
|
|
|
|
def p_structs(self, structs: list[Struct]):
|
|
for struct in structs:
|
|
self._p_line(lambda: self.p_struct(struct))
|
|
|
|
def p_struct(self, struct: Struct):
|
|
self._p_comment(struct.description, doc=True)
|
|
self._p_line(lambda: self._p_str(f"struct {struct.name} {{"))
|
|
self._with_indent(lambda: self.p_struct_fields(struct.fields))
|
|
self._p_line(lambda: self._p_str("}"))
|
|
|
|
def p_struct_fields(self, fields: list[StructField]):
|
|
for field in fields:
|
|
self._p_line(lambda: self.p_struct_field(field))
|
|
|
|
def p_struct_field(self, field: StructField):
|
|
self._p_comment(field.description)
|
|
self._p_indented(lambda: self._p_str(f"{field.ty} {field.name};"))
|
|
|
|
def p_functions(self, cheatcodes: list[Cheatcode]):
|
|
for cheatcode in cheatcodes:
|
|
self._p_line(lambda: self.p_function(cheatcode.func))
|
|
|
|
def p_function(self, func: Function):
|
|
self._p_comment(func.description, doc=True)
|
|
self._p_line(lambda: self._p_str(func.declaration))
|
|
|
|
def _p_comment(self, s: str, doc: bool = False):
|
|
s = s.strip()
|
|
if s == "":
|
|
return
|
|
|
|
s = map(lambda line: line.lstrip(), s.split("\n"))
|
|
if self.block_doc_style:
|
|
self._p_str("/*")
|
|
if doc:
|
|
self._p_str("*")
|
|
self._p_nl()
|
|
for line in s:
|
|
self._p_indent()
|
|
self._p_str(" ")
|
|
if doc:
|
|
self._p_str("* ")
|
|
self._p_str(line)
|
|
self._p_nl()
|
|
self._p_indent()
|
|
self._p_str(" */")
|
|
self._p_nl()
|
|
else:
|
|
first_line = True
|
|
for line in s:
|
|
if not first_line:
|
|
self._p_indent()
|
|
first_line = False
|
|
|
|
if doc:
|
|
self._p_str("/// ")
|
|
else:
|
|
self._p_str("// ")
|
|
self._p_str(line)
|
|
self._p_nl()
|
|
|
|
def _with_indent(self, f: VoidFn):
|
|
self._inc_indent()
|
|
f()
|
|
self._dec_indent()
|
|
|
|
def _p_line(self, f: VoidFn):
|
|
self._p_indent()
|
|
f()
|
|
self._p_nl()
|
|
|
|
def _p_indented(self, f: VoidFn):
|
|
self._p_indent()
|
|
f()
|
|
|
|
def _p_indent(self):
|
|
for _ in range(self.indent_level):
|
|
self._p_str(self._indent_str)
|
|
|
|
def _p_nl(self):
|
|
self._p_str(self.nl_str)
|
|
|
|
def _p_str(self, txt: str):
|
|
self.buffer += txt
|
|
|
|
def _inc_indent(self):
|
|
self.indent_level += 1
|
|
|
|
def _dec_indent(self):
|
|
self.indent_level -= 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|