feat[venom]: add IRContext prefix + merge

tests/unit/compiler/venom/test_context_prefix.py

@@ -0,0 +1,155 @@
+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").named_label("foo").value == "m1_foo"
+    assert IRContext().named_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))

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.named_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.named_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.named_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))

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.named_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.named_label(src)
         self._emit_evm("codecopy", dst, src, size)
 
     def extcodecopy(self, addr: Operand, dst: IRVariable, src: Operand, size: Operand) -> None: