Skip to content

Commit 8a8f35b

Browse files
charles-coopercharles-cooper
authored
Merge branch 'master' into fix/top-level-abi-key
2 parents 91566d3 + 4845fd4 commit 8a8f35b

File tree

8 files changed

+152
-37
lines changed

8 files changed

+152
-37
lines changed

‎tests/unit/ast/test_ast_dict.py‎

Lines changed: 54 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import copy
12
import json
23

34
from vyper import compiler
@@ -216,24 +217,27 @@ def foo():
216217
input_bundle = make_input_bundle({"lib1.vy": lib1, "main.vy": main})
217218

218219
lib1_file = input_bundle.load_file("lib1.vy")
219-
out = compiler.compile_from_file_input(
220+
lib1_out = compiler.compile_from_file_input(
220221
lib1_file, input_bundle=input_bundle, output_formats=["annotated_ast_dict"]
221222
)
222-
lib1_ast = out["annotated_ast_dict"]["ast"]
223+
224+
lib1_ast = copy.deepcopy(lib1_out["annotated_ast_dict"]["ast"])
223225
lib1_sha256sum = lib1_ast.pop("source_sha256sum")
224226
assert lib1_sha256sum == lib1_file.sha256sum
225227
to_strip = NODE_SRC_ATTRIBUTES + ("resolved_path", "variable_reads", "variable_writes")
226228
_strip_source_annotations(lib1_ast, to_strip=to_strip)
227229

228230
main_file = input_bundle.load_file("main.vy")
229-
out = compiler.compile_from_file_input(
231+
main_out = compiler.compile_from_file_input(
230232
main_file, input_bundle=input_bundle, output_formats=["annotated_ast_dict"]
231233
)
232-
main_ast = out["annotated_ast_dict"]["ast"]
234+
main_ast = main_out["annotated_ast_dict"]["ast"]
233235
main_sha256sum = main_ast.pop("source_sha256sum")
234236
assert main_sha256sum == main_file.sha256sum
235237
_strip_source_annotations(main_ast, to_strip=to_strip)
236238

239+
assert main_out["annotated_ast_dict"]["imports"][0] == lib1_out["annotated_ast_dict"]["ast"]
240+
237241
# TODO: would be nice to refactor this into bunch of small test cases
238242
assert main_ast == {
239243
"ast_type": "Module",
@@ -1776,3 +1780,49 @@ def qux2():
17761780
},
17771781
}
17781782
]
1783+
1784+
1785+
def test_annotated_ast_export_recursion(make_input_bundle):
1786+
sources = {
1787+
"main.vy": """
1788+
import lib1
1789+
1790+
@external
1791+
def foo():
1792+
lib1.foo()
1793+
""",
1794+
"lib1.vy": """
1795+
import lib2
1796+
1797+
def foo():
1798+
lib2.foo()
1799+
""",
1800+
"lib2.vy": """
1801+
def foo():
1802+
pass
1803+
""",
1804+
}
1805+
1806+
input_bundle = make_input_bundle(sources)
1807+
1808+
def compile_and_get_ast(file_name):
1809+
file = input_bundle.load_file(file_name)
1810+
output = compiler.compile_from_file_input(
1811+
file, input_bundle=input_bundle, output_formats=["annotated_ast_dict"]
1812+
)
1813+
return output["annotated_ast_dict"]
1814+
1815+
lib1_ast = compile_and_get_ast("lib1.vy")["ast"]
1816+
lib2_ast = compile_and_get_ast("lib2.vy")["ast"]
1817+
main_out = compile_and_get_ast("main.vy")
1818+
1819+
lib1_import_ast = main_out["imports"][1]
1820+
lib2_import_ast = main_out["imports"][0]
1821+
1822+
# path is once virtual, once libX.vy
1823+
# type contains name which is based on path
1824+
keys = [s for s in lib1_import_ast.keys() if s not in {"path", "type"}]
1825+
1826+
for key in keys:
1827+
assert lib1_ast[key] == lib1_import_ast[key]
1828+
assert lib2_ast[key] == lib2_import_ast[key]

‎vyper/ast/nodes.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -949,6 +949,12 @@ def validate(self):
949949
class Ellipsis(Constant):
950950
__slots__ = ()
951951

952+
def to_dict(self):
953+
ast_dict = super().to_dict()
954+
# python ast ellipsis() is not json serializable; use a string
955+
ast_dict["value"] = self.node_source_code
956+
return ast_dict
957+
952958

953959
class Dict(ExprNode):
954960
__slots__ = ("keys", "values")

‎vyper/ast/nodes.pyi‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@ class TopLevel(VyperNode):
7070
class Module(TopLevel):
7171
path: str = ...
7272
resolved_path: str = ...
73+
source_id: int = ...
7374
def namespace(self) -> Any: ... # context manager
7475

7576
class FunctionDef(TopLevel):

‎vyper/compiler/output.py‎

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,18 @@
33
from collections import deque
44
from pathlib import PurePath
55

6-
from vyper.ast import ast_to_dict
6+
import vyper.ast as vy_ast
7+
from vyper.ast.utils import ast_to_dict
78
from vyper.codegen.ir_node import IRnode
89
from vyper.compiler.output_bundle import SolcJSONWriter, VyperArchiveWriter
910
from vyper.compiler.phases import CompilerData
1011
from vyper.compiler.utils import build_gas_estimates
1112
from vyper.evm import opcodes
1213
from vyper.exceptions import VyperException
1314
from vyper.ir import compile_ir
15+
from vyper.semantics.analysis.base import ModuleInfo
1416
from vyper.semantics.types.function import FunctionVisibility, StateMutability
17+
from vyper.semantics.types.module import InterfaceT
1518
from vyper.typing import StorageLayout
1619
from vyper.utils import vyper_warn
1720
from vyper.warnings import ContractSizeLimitWarning
@@ -26,9 +29,32 @@ def build_ast_dict(compiler_data: CompilerData) -> dict:
2629

2730

2831
def build_annotated_ast_dict(compiler_data: CompilerData) -> dict:
32+
module_t = compiler_data.annotated_vyper_module._metadata["type"]
33+
# get all reachable imports including recursion
34+
imported_module_infos = module_t.reachable_imports
35+
unique_modules: dict[str, vy_ast.Module] = {}
36+
for info in imported_module_infos:
37+
if isinstance(info.typ, InterfaceT):
38+
ast = info.typ.decl_node
39+
if ast is None: # json abi
40+
continue
41+
else:
42+
assert isinstance(info.typ, ModuleInfo)
43+
ast = info.typ.module_t._module
44+
45+
assert isinstance(ast, vy_ast.Module) # help mypy
46+
# use resolved_path for uniqueness, since Module objects can actually
47+
# come from multiple InputBundles (particularly builtin interfaces),
48+
# so source_id is not guaranteed to be unique.
49+
if ast.resolved_path in unique_modules:
50+
# sanity check -- objects must be identical
51+
assert unique_modules[ast.resolved_path] is ast
52+
unique_modules[ast.resolved_path] = ast
53+
2954
annotated_ast_dict = {
3055
"contract_name": str(compiler_data.contract_path),
3156
"ast": ast_to_dict(compiler_data.annotated_vyper_module),
57+
"imports": [ast_to_dict(ast) for ast in unique_modules.values()],
3258
}
3359
return annotated_ast_dict
3460

‎vyper/venom/ir_node_to_venom.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -548,7 +548,7 @@ def emit_body_blocks():
548548
_global_symbols[ir.value] = ptr
549549
elif ir.value.startswith("$palloca") and ir.value not in _global_symbols:
550550
alloca = ir.passthrough_metadata["alloca"]
551-
ptr = fn.get_basic_block().append_instruction("store", alloca.offset)
551+
ptr = fn.get_basic_block().append_instruction("palloca", alloca.offset, alloca.size)
552552
_global_symbols[ir.value] = ptr
553553

554554
return _global_symbols.get(ir.value) or symbols.get(ir.value)

‎vyper/venom/passes/mem2var.py‎

Lines changed: 38 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,7 @@
1-
from vyper.utils import OrderedSet
21
from vyper.venom.analysis.cfg import CFGAnalysis
32
from vyper.venom.analysis.dfg import DFGAnalysis
43
from vyper.venom.analysis.liveness import LivenessAnalysis
5-
from vyper.venom.basicblock import IRBasicBlock, IRInstruction, IRVariable
4+
from vyper.venom.basicblock import IRInstruction, IRVariable
65
from vyper.venom.function import IRFunction
76
from vyper.venom.passes.base_pass import IRPass
87

@@ -14,21 +13,27 @@ class Mem2Var(IRPass):
1413
"""
1514

1615
function: IRFunction
17-
defs: dict[IRVariable, OrderedSet[IRBasicBlock]]
1816

1917
def run_pass(self):
2018
self.analyses_cache.request_analysis(CFGAnalysis)
2119
dfg = self.analyses_cache.request_analysis(DFGAnalysis)
2220

2321
self.var_name_count = 0
2422
for var, inst in dfg.outputs.items():
25-
if inst.opcode != "alloca":
26-
continue
27-
self._process_alloca_var(dfg, var)
23+
if inst.opcode == "alloca":
24+
self._process_alloca_var(dfg, var)
25+
elif inst.opcode == "palloca":
26+
self._process_palloca_var(dfg, inst, var)
2827

2928
self.analyses_cache.invalidate_analysis(DFGAnalysis)
3029
self.analyses_cache.invalidate_analysis(LivenessAnalysis)
3130

31+
def _mk_varname(self, varname: str):
32+
varname = varname.removeprefix("%")
33+
varname = f"var{varname}_{self.var_name_count}"
34+
self.var_name_count += 1
35+
return varname
36+
3237
def _process_alloca_var(self, dfg: DFGAnalysis, var: IRVariable):
3338
"""
3439
Process alloca allocated variable. If it is only used by mstore/mload/return
@@ -40,8 +45,7 @@ def _process_alloca_var(self, dfg: DFGAnalysis, var: IRVariable):
4045
elif all([inst.opcode == "mstore" for inst in uses]):
4146
return
4247
elif all([inst.opcode in ["mstore", "mload", "return"] for inst in uses]):
43-
var_name = f"addr{var.name}_{self.var_name_count}"
44-
self.var_name_count += 1
48+
var_name = self._mk_varname(var.name)
4549
for inst in uses:
4650
if inst.opcode == "mstore":
4751
inst.opcode = "store"
@@ -52,7 +56,32 @@ def _process_alloca_var(self, dfg: DFGAnalysis, var: IRVariable):
5256
inst.operands = [IRVariable(var_name)]
5357
elif inst.opcode == "return":
5458
bb = inst.parent
55-
idx = bb.instructions.index(inst)
59+
idx = len(bb.instructions) - 1
60+
assert inst == bb.instructions[idx] # sanity
5661
bb.insert_instruction(
5762
IRInstruction("mstore", [IRVariable(var_name), inst.operands[1]]), idx
5863
)
64+
65+
def _process_palloca_var(self, dfg: DFGAnalysis, palloca_inst: IRInstruction, var: IRVariable):
66+
"""
67+
Process alloca allocated variable. If it is only used by mstore/mload
68+
instructions, it is promoted to a stack variable. Otherwise, it is left as is.
69+
"""
70+
uses = dfg.get_uses(var)
71+
if not all(inst.opcode in ["mstore", "mload"] for inst in uses):
72+
return
73+
74+
var_name = self._mk_varname(var.name)
75+
76+
palloca_inst.opcode = "mload"
77+
palloca_inst.operands = [palloca_inst.operands[0]]
78+
palloca_inst.output = IRVariable(var_name)
79+
80+
for inst in uses:
81+
if inst.opcode == "mstore":
82+
inst.opcode = "store"
83+
inst.output = IRVariable(var_name)
84+
inst.operands = [inst.operands[0]]
85+
elif inst.opcode == "mload":
86+
inst.opcode = "store"
87+
inst.operands = [IRVariable(var_name)]

‎vyper/venom/passes/sccp/sccp.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -179,7 +179,7 @@ def _visit_phi(self, inst: IRInstruction):
179179

180180
def _visit_expr(self, inst: IRInstruction):
181181
opcode = inst.opcode
182-
if opcode in ["store", "alloca"]:
182+
if opcode in ["store", "alloca", "palloca"]:
183183
assert inst.output is not None, "Got store/alloca without output"
184184
out = self._eval_from_lattice(inst.operands[0])
185185
self._set_lattice(inst.output, out)

‎vyper/venom/venom_to_assembly.py‎

Lines changed: 24 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -295,7 +295,8 @@ def _generate_evm_for_basicblock_r(
295295
asm.append(f"_sym_{basicblock.label}")
296296
asm.append("JUMPDEST")
297297

298-
self.clean_stack_from_cfg_in(asm, basicblock, stack)
298+
if len(basicblock.cfg_in) == 1:
299+
self.clean_stack_from_cfg_in(asm, basicblock, stack)
299300

300301
all_insts = sorted(basicblock.instructions, key=lambda x: x.opcode != "param")
301302

@@ -321,26 +322,28 @@ def _generate_evm_for_basicblock_r(
321322
def clean_stack_from_cfg_in(
322323
self, asm: list, basicblock: IRBasicBlock, stack: StackModel
323324
) -> None:
324-
if len(basicblock.cfg_in) == 0:
325-
return
326-
327-
to_pop = OrderedSet[IRVariable]()
328-
for in_bb in basicblock.cfg_in:
329-
# inputs is the input variables we need from in_bb
330-
inputs = self.liveness_analysis.input_vars_from(in_bb, basicblock)
331-
332-
# layout is the output stack layout for in_bb (which works
333-
# for all possible cfg_outs from the in_bb).
334-
layout = in_bb.out_vars
335-
336-
# pop all the stack items which in_bb produced which we don't need.
337-
to_pop |= layout.difference(inputs)
338-
325+
# the input block is a splitter block, like jnz or djmp
326+
assert len(basicblock.cfg_in) == 1
327+
in_bb = basicblock.cfg_in.first()
328+
assert len(in_bb.cfg_out) > 1
329+
330+
# inputs is the input variables we need from in_bb
331+
inputs = self.liveness_analysis.input_vars_from(in_bb, basicblock)
332+
333+
# layout is the output stack layout for in_bb (which works
334+
# for all possible cfg_outs from the in_bb, in_bb is responsible
335+
# for making sure its output stack layout works no matter which
336+
# bb it jumps into).
337+
layout = in_bb.out_vars
338+
to_pop = list(layout.difference(inputs))
339+
340+
# small heuristic: pop from shallowest first.
341+
to_pop.sort(key=lambda var: -stack.get_depth(var))
342+
343+
# NOTE: we could get more fancy and try to optimize the swap
344+
# operations here, there is probably some more room for optimization.
339345
for var in to_pop:
340346
depth = stack.get_depth(var)
341-
# don't pop phantom phi inputs
342-
if depth is StackModel.NOT_IN_STACK:
343-
continue
344347

345348
if depth != 0:
346349
self.swap(asm, stack, depth)
@@ -360,7 +363,7 @@ def _generate_evm_for_instruction(
360363

361364
if opcode in ["jmp", "djmp", "jnz", "invoke"]:
362365
operands = list(inst.get_non_label_operands())
363-
elif opcode == "alloca":
366+
elif opcode in ("alloca", "palloca"):
364367
offset, _size = inst.operands
365368
operands = [offset]
366369

@@ -460,7 +463,7 @@ def _generate_evm_for_instruction(
460463
# Step 5: Emit the EVM instruction(s)
461464
if opcode in _ONE_TO_ONE_INSTRUCTIONS:
462465
assembly.append(opcode.upper())
463-
elif opcode == "alloca":
466+
elif opcode in ("alloca", "palloca"):
464467
pass
465468
elif opcode == "param":
466469
pass

0 commit comments

Comments
 (0)