diff --git a/tests/unit/compiler/venom/test_context_prefix.py b/tests/unit/compiler/venom/test_context_prefix.py new file mode 100644 index 0000000000..7b24cfcb6e --- /dev/null +++ b/tests/unit/compiler/venom/test_context_prefix.py @@ -0,0 +1,163 @@ +import pytest + +from vyper.venom.basicblock import IRBasicBlock, IRLabel +from vyper.venom.context import IRContext + + +def _fn_labels(ctx: IRContext) -> list[str]: + return sorted(label.value for label in ctx.functions) + + +def _section_labels(ctx: IRContext) -> list[str]: + return [section.label.value for section in ctx.data_segment] + + +@pytest.mark.parametrize( + "prefix, expected_fn, expected_section, expected_label", + [("", "foo", "tbl", "1"), ("m1", "m1_foo", "m1_tbl", "m1_1")], +) +def test_prefix_applies_to_generated_names(prefix, expected_fn, expected_section, expected_label): + ctx = IRContext(prefix=prefix) + fn = ctx.create_function("foo") + ctx.append_data_section("tbl") + label = ctx.get_next_label() + + assert fn.name.value == expected_fn + assert _section_labels(ctx) == [expected_section] + assert label.value == expected_label + + +def test_prefix_applies_to_labeled_helpers(): + assert IRContext(prefix="m1").prefixed_label("foo").value == "m1_foo" + assert IRContext().prefixed_label("foo").value == "foo" + + +def test_explicit_irlabel_passes_through(): + # IRLabel arg is used as-is; only str input to append_data_section is auto-prefixed. + ctx = IRContext(prefix="m1") + ctx.append_data_section(IRLabel("user_built", is_symbol=True)) + assert _section_labels(ctx) == ["user_built"] + + +def test_prefix_applies_to_suffixed_labels(): + ctx = IRContext(prefix="m1") + assert ctx.get_next_label("loop").value == "m1_1_loop" + + +def test_merge_moves_state_and_clears_sources(): + a = IRContext(prefix="a") + a.entry_function = a.create_function("foo") + a.append_data_section("v") + + b = IRContext(prefix="b") + b.create_function("bar") + b.append_data_section("w") + + target = IRContext() + assert target.merge(a, b) is target + + assert _fn_labels(target) == ["a_foo", "b_bar"] + assert _section_labels(target) == ["a_v", "b_w"] + + assert a.functions == {} + assert a.data_segment == [] + assert a.entry_function is None + assert b.functions == {} + assert b.data_segment == [] + + +@pytest.mark.parametrize( + "target_prefix, src1_prefix, src2_prefix, expected_message", + [("", "dup", "dup", "duplicate function"), ("t", "t", None, "duplicate function")], +) +def test_merge_raises_on_duplicate_function_labels( + target_prefix, src1_prefix, src2_prefix, expected_message +): + target = IRContext(prefix=target_prefix) + if target_prefix: + target.create_function("foo") + + src1 = IRContext(prefix=src1_prefix) + src1.create_function("foo") + + with pytest.raises(ValueError, match=expected_message): + if src2_prefix is None: + target.merge(src1) + else: + src2 = IRContext(prefix=src2_prefix) + src2.create_function("foo") + target.merge(src1, src2) + + +def test_merge_raises_on_duplicate_basic_block_labels(): + # Distinct function names but a shared prefix → bb labels from get_next_label collide. + a = IRContext(prefix="m") + fn_a = a.create_function("foo") + fn_a.append_basic_block(IRBasicBlock(a.get_next_label(), fn_a)) + + b = IRContext(prefix="m") + fn_b = b.create_function("bar") + fn_b.append_basic_block(IRBasicBlock(b.get_next_label(), fn_b)) + + with pytest.raises(ValueError, match="duplicate basic block label"): + IRContext().merge(a, b) + + +def test_merge_raises_on_duplicate_data_section_label(): + a = IRContext(prefix="dup") + b = IRContext(prefix="dup") + a.append_data_section("tbl") + b.append_data_section("tbl") + + with pytest.raises(ValueError, match="duplicate data section"): + IRContext().merge(a, b) + + +def test_merge_is_atomic_on_validation_failure(): + target = IRContext(prefix="t") + target.create_function("target") + target.append_data_section("target_tbl") + + src_ok = IRContext(prefix="good") + src_ok.create_function("foo") + src_ok.append_data_section("tbl") + + src_bad = IRContext(prefix="good") + src_bad.create_function("foo") + src_bad.append_data_section("tbl") + + with pytest.raises(ValueError, match="duplicate function"): + target.merge(src_ok, src_bad) + + assert _fn_labels(target) == ["t_target"] + assert _section_labels(target) == ["t_target_tbl"] + assert _fn_labels(src_ok) == ["good_foo"] + assert _section_labels(src_ok) == ["good_tbl"] + assert _fn_labels(src_bad) == ["good_foo"] + assert _section_labels(src_bad) == ["good_tbl"] + + +def test_prefixed_labels_roundtrip_through_parser(): + from vyper.venom.parser import parse_venom + + ctx = IRContext(prefix="m1") + fn = ctx.create_function("foo") + extra = IRBasicBlock(ctx.get_next_label("loop"), fn) + fn.append_basic_block(extra) + fn.entry.append_instruction("jmp", extra.label) + extra.append_instruction("stop") + + parse_venom(str(ctx)) + + +def test_merge_advances_counters_past_sources(): + src = IRContext(prefix="m") + fn = src.create_function("foo") + fn.append_basic_block(IRBasicBlock(src.get_next_label(), fn)) # "m_1" + fn.append_basic_block(IRBasicBlock(src.get_next_label(), fn)) # "m_2" + src.get_next_variable() # "%1" + src.get_next_variable() # "%2" + target = IRContext(prefix="m") + target.merge(src) + assert target.get_next_label().value == "m_3" + assert target.get_next_variable().value == "%3" diff --git a/vyper/codegen_venom/module.py b/vyper/codegen_venom/module.py index 7e1224626c..4f37a2f8d2 100644 --- a/vyper/codegen_venom/module.py +++ b/vyper/codegen_venom/module.py @@ -182,12 +182,12 @@ def generate_deploy_venom( deploy_ctx = IRContext() # Add runtime bytecode as data section - deploy_ctx.append_data_section(IRLabel("runtime_begin")) + deploy_ctx.append_data_section("runtime_begin") deploy_ctx.append_data_item(runtime_bytecode) # Add CBOR metadata if provided if cbor_metadata is not None: - deploy_ctx.append_data_section(IRLabel("cbor_metadata")) + deploy_ctx.append_data_section("cbor_metadata") deploy_ctx.append_data_item(cbor_metadata) deploy_fn = deploy_ctx.create_function("deploy") @@ -431,13 +431,13 @@ def _generate_selector_section_sparse( bucket_id = builder.mod(method_id, IRLiteral(n_buckets)) # Create data section with bucket headers - runtime_ctx.append_data_section(IRLabel("selector_buckets", is_symbol=True)) + runtime_ctx.append_data_section("selector_buckets") # Build jump targets list and add bucket header labels jump_targets = [] for i in range(n_buckets): if i in buckets: - bucket_label = IRLabel(f"selector_bucket_{i}", is_symbol=True) + bucket_label = runtime_ctx.prefixed_label(f"selector_bucket_{i}") jump_targets.append(bucket_label) else: # Empty bucket -> fallback @@ -448,9 +448,7 @@ def _generate_selector_section_sparse( # Location = selector_buckets + bucket_id * 2 bucket_hdr_offset = builder.mul(bucket_id, IRLiteral(SZ_BUCKET_HEADER)) # Use add with label - the label resolves to its code position at link time - selector_buckets_addr = builder.offset( - IRLiteral(0), IRLabel("selector_buckets", is_symbol=True) - ) + selector_buckets_addr = builder.offset(IRLiteral(0), "selector_buckets") bucket_hdr_location = builder.add(selector_buckets_addr, bucket_hdr_offset) # Copy 2-byte header to memory at offset (32 - 2) = 30 @@ -468,7 +466,7 @@ def _generate_selector_section_sparse( # Generate bucket blocks for bucket_id_val, bucket_method_ids in buckets.items(): - bucket_label = IRLabel(f"selector_bucket_{bucket_id_val}", is_symbol=True) + bucket_label = runtime_ctx.prefixed_label(f"selector_bucket_{bucket_id_val}") bucket_bb = builder.create_block(f"bucket_{bucket_id_val}") # Override the label to match the data section reference bucket_bb.label = bucket_label @@ -695,25 +693,23 @@ def _generate_selector_section_dense( entry_point_labels: dict[str, IRLabel] = {} for abi_sig, (_func_ast, _entry_info) in all_entry_points.items(): method_id_val = method_id_int(abi_sig) - label = IRLabel(f"entry_{method_id_val:08x}", is_symbol=True) + label = runtime_ctx.prefixed_label(f"entry_{method_id_val:08x}") entry_point_labels[abi_sig] = label # Compute bucket_id = method_id % n_buckets bucket_id_var = builder.mod(method_id, IRLiteral(n_buckets)) # Create data section for bucket headers - runtime_ctx.append_data_section(IRLabel("BUCKET_HEADERS", is_symbol=True)) + runtime_ctx.append_data_section("BUCKET_HEADERS") for bucket_id_val, bucket in sorted(jumptable_info.items()): runtime_ctx.append_data_item(bucket.magic.to_bytes(2, "big")) - runtime_ctx.append_data_item(IRLabel(f"bucket_{bucket_id_val}", is_symbol=True)) + runtime_ctx.append_data_item(f"bucket_{bucket_id_val}") runtime_ctx.append_data_item(bucket.bucket_size.to_bytes(1, "big")) # Load bucket header from data section # Location = BUCKET_HEADERS + bucket_id * 5 bucket_hdr_offset = builder.mul(bucket_id_var, IRLiteral(SZ_BUCKET_HEADER)) - bucket_headers_addr = builder.offset( - IRLiteral(0), IRLabel("BUCKET_HEADERS", is_symbol=True) - ) + bucket_headers_addr = builder.offset(IRLiteral(0), "BUCKET_HEADERS") bucket_hdr_location = builder.add(bucket_headers_addr, bucket_hdr_offset) # Copy 5-byte header to memory at offset (32 - 5) = 27 @@ -798,7 +794,7 @@ def _generate_selector_section_dense( # Create data sections for each bucket's function info for bucket_id_val, bucket in jumptable_info.items(): - runtime_ctx.append_data_section(IRLabel(f"bucket_{bucket_id_val}", is_symbol=True)) + runtime_ctx.append_data_section(f"bucket_{bucket_id_val}") # Sort function infos by their image (hash position) for mid in bucket.method_ids_image_order: @@ -1523,7 +1519,7 @@ def _emit_deploy_epilogue( builder.assert_(copy_success) # Copy runtime bytecode from data section to memory - builder.codecopy(dst_ptr, IRLabel("runtime_begin"), IRLiteral(runtime_codesize)) + builder.codecopy(dst_ptr, "runtime_begin", IRLiteral(runtime_codesize)) # Return runtime + immutables builder.return_(dst_ptr, IRLiteral(total_size)) diff --git a/vyper/venom/builder.py b/vyper/venom/builder.py index ce54cf92da..39304af969 100644 --- a/vyper/venom/builder.py +++ b/vyper/venom/builder.py @@ -277,11 +277,13 @@ def istore(self, offset: IRVariable, val: Operand) -> None: """Store val to immutable memory region at offset (deploy-time only). (IR-specific)""" self._emit("istore", offset, val) - def offset(self, operand: Operand, label: IRLabel) -> IRVariable: + def offset(self, operand: Operand, label: IRLabel | str) -> IRVariable: """Compute static offset from label. Used for code position calculations. (IR-specific) Computes label + operand. Args order matches Venom IR: offset operand, @label """ + if isinstance(label, str): + label = self.ctx.prefixed_label(label) return self._emit1("offset", operand, label) # === Control Flow (IR-specific) === @@ -394,8 +396,10 @@ def calldatacopy(self, dst: IRVariable, src: Operand, size: Operand) -> None: """Copy size bytes from calldata[src] to memory[dst].""" self._emit_evm("calldatacopy", dst, src, size) - def codecopy(self, dst: IRVariable, src: Operand, size: Operand) -> None: + def codecopy(self, dst: IRVariable, src: Operand | str, size: Operand) -> None: """Copy size bytes from code[src] to memory[dst].""" + if isinstance(src, str): + src = self.ctx.prefixed_label(src) self._emit_evm("codecopy", dst, src, size) def extcodecopy(self, addr: Operand, dst: IRVariable, src: Operand, size: Operand) -> None: diff --git a/vyper/venom/context.py b/vyper/venom/context.py index 5842fcb63d..d19974f5d2 100644 --- a/vyper/venom/context.py +++ b/vyper/venom/context.py @@ -48,8 +48,9 @@ class IRContext: last_variable: int mem_allocator: MemoryAllocator global_analyses_cache: Optional["IRGlobalAnalysesCache"] + prefix: str - def __init__(self) -> None: + def __init__(self, prefix: str = "") -> None: self.functions = {} self.entry_function = None self.data_segment = [] @@ -59,22 +60,32 @@ def __init__(self) -> None: self.mem_allocator = MemoryAllocator() self.global_analyses_cache = None + self.prefix = prefix def get_basic_blocks(self) -> Iterator[IRBasicBlock]: for fn in self.functions.values(): for bb in fn.get_basic_blocks(): yield bb + def _prefixed_value(self, value: str) -> str: + return f"{self.prefix}_{value}" if self.prefix else value + def add_function(self, fn: IRFunction) -> None: + assert fn.name not in self.functions, f"duplicate function {fn.name}" fn.ctx = self self.functions[fn.name] = fn def remove_function(self, fn: IRFunction) -> None: del self.functions[fn.name] + def prefixed_label(self, name: str, is_symbol: bool = True) -> IRLabel: + """Return ``IRLabel(f"{prefix}_{name}")`` (or ``IRLabel(name)`` if + prefix is empty). Use for labels that must survive a :meth:`merge`. + """ + return IRLabel(self._prefixed_value(name), is_symbol=is_symbol) + def create_function(self, name: str) -> IRFunction: - label = IRLabel(name, True) - assert label not in self.functions, f"duplicate function {label}" + label = self.prefixed_label(name, is_symbol=True) fn = IRFunction(label, self) self.add_function(fn) return fn @@ -88,10 +99,9 @@ def get_functions(self) -> Iterator[IRFunction]: return iter(self.functions.values()) def get_next_label(self, suffix: str = "") -> IRLabel: - if suffix != "": - suffix = f"_{suffix}" + suffix = f"_{suffix}" if suffix else "" self.last_label += 1 - return IRLabel(f"{self.last_label}{suffix}") + return IRLabel(self._prefixed_value(f"{self.last_label}{suffix}")) def get_next_variable(self) -> IRVariable: self.last_variable += 1 @@ -100,13 +110,62 @@ def get_next_variable(self) -> IRVariable: def get_last_variable(self) -> str: return f"%{self.last_variable}" - def append_data_section(self, name: IRLabel) -> None: + def append_data_section(self, name: IRLabel | str) -> None: + """``str`` → auto-namespaced via :meth:`prefixed_label`; ``IRLabel`` → used as-is.""" + if isinstance(name, str): + name = self.prefixed_label(name) self.data_segment.append(DataSection(name)) - def append_data_item(self, data: IRLabel | bytes) -> None: + def merge(self, *sources: "IRContext") -> "IRContext": + """Splice each source's functions / data sections into ``self``; clears + the sources. Raises :class:`ValueError` on label clash before mutating. + """ + function_labels = set(self.functions) + data_labels = {section.label for section in self.data_segment} + bb_labels = {bb.label for bb in self.get_basic_blocks()} + + for src in sources: + for fn in src.functions.values(): + if fn.name in function_labels: + raise ValueError( + f"merge: duplicate function label {fn.name}; " + "two sources share a prefix or collide with the target" + ) + function_labels.add(fn.name) + + for bb in fn.get_basic_blocks(): + if bb.label in bb_labels: + raise ValueError( + f"merge: duplicate basic block label {bb.label}; " + "two sources share a prefix or collide with the target" + ) + bb_labels.add(bb.label) + + for section in src.data_segment: + if section.label in data_labels: + raise ValueError( + f"merge: duplicate data section label {section.label}; " + "two sources share a prefix or collide with the target" + ) + data_labels.add(section.label) + + for src in sources: + for fn in list(src.functions.values()): + self.add_function(fn) + self.data_segment.extend(src.data_segment) + self.last_label = max(self.last_label, src.last_label) + self.last_variable = max(self.last_variable, src.last_variable) + src.functions.clear() + src.data_segment.clear() + src.entry_function = None + return self + + def append_data_item(self, data: IRLabel | bytes | str) -> None: """ Append data to current data section """ + if isinstance(data, str): + data = self.prefixed_label(data) assert len(self.data_segment) > 0 data_section = self.data_segment[-1] data_section.data_items.append(DataItem(data))