Code Section Specification

magia/module.py

@@ -4,7 +4,7 @@
 import logging
 from collections import Counter, OrderedDict
 from dataclasses import dataclass
-from functools import cached_property
+from functools import cached_property, partial
 from itertools import count
 from os import PathLike
 from string import Template
@@ -13,7 +13,7 @@
 from .data_struct import SignalDict, SignalType
 from .io_ports import IOPorts
 from .memory import Memory, MemorySignal
-from .signals import SIGNAL_ASSIGN_TEMPLATE, Signal, Synthesizable
+from .signals import SIGNAL_ASSIGN_TEMPLATE, CodeSectionType, Signal, Synthesizable
 
 if TYPE_CHECKING:
     from .bundle import Bundle
@@ -38,10 +38,13 @@ class ModuleInstanceConfig:
 IO_TEMPLATE = Template(".$port_name($signal_name)")
 
 
+class _ModuleDefaultCodeSection(type):
+    def __call__(cls, *args, **kwargs):
+        with Synthesizable.code_section(CodeSectionType.LOGIC):
+            return super().__call__(*args, **kwargs)
 
 
-
-class Module(Synthesizable):
+class Module(Synthesizable, metaclass=_ModuleDefaultCodeSection):
     """
     A module is a collection of signals and operations. It can also include other modules.
 
@@ -69,6 +72,8 @@ def implement(self):
     _new_module_counter = count(0)
     output_file: None | PathLike = None
 
+    formal_code = partial(Synthesizable.code_section, CodeSectionType.FORMAL)
+
     def __init__(self, name: None | str = None, **kwargs):
         super().__init__(**kwargs)
 
@@ -461,13 +466,10 @@ def __init__(self, module: Module, **kwargs):
         if kwargs.get("name") is None and module.name is not None:
             kwargs["name"] = f"{module.name}Wrapper"
         super().__init__(**kwargs)
-        self.module = module
-        self.io += module.io
-        module.instance(name="inst", io={
-            name: self.io[name]
-            for name in module.io.signals
-        })
-
-    def elaborate(self) -> tuple[str, set[Module]]:
-        with Signal.decl_in_verilog():
-            return super().elaborate()
+        with self.code_section(CodeSectionType.VERILOG):
+            self.module = module
+            self.io += module.io
+            module.instance(name="inst", io={
+                name: self.io[name]
+                for name in module.io.signals
+            })

magia/signals.py

@@ -4,6 +4,7 @@
 from collections.abc import Iterable
 from contextlib import contextmanager
 from dataclasses import dataclass
+from enum import Enum, auto
 from functools import cached_property
 from itertools import count
 from pathlib import Path
@@ -22,10 +23,17 @@
 CURRENT_DIR = Path(__file__).parent
 
 SIGNAL_DECL_TEMPLATE = Template("logic $signed $width $name;")
+SIGNAL_DECL_FORMAL_TEMPLATE = Template("logic $signed $width $name = 0;")
 SIGNAL_DECL_VERILOG_TEMPLATE = Template("wire $signed $width $name;")
 SIGNAL_ASSIGN_TEMPLATE = Template("assign $name = $driver;")
 
 
+class CodeSectionType(Enum):
+    LOGIC = auto()
+    VERILOG = auto()
+    FORMAL = auto()
+
+
 @dataclass
 class SignalConfig:
     name: None | str = None
@@ -55,13 +63,30 @@ class Synthesizable:
 
     _ANNOTATION_TEMPLATE = Template("/*\nNet name: $net_name\n$comment$loc\n*/")
 
+    current_code_section = CodeSectionType.LOGIC
+
+    @classmethod
+    @contextmanager
+    def code_section(cls, section: CodeSectionType):
+        """
+        Specify the code section of the synthesizable objects created within this context manager.
+
+        Code section specify how the object is elaborated in the SystemVerilog code.
+
+        :param section: The type of code section.
+        """
+        prev_value, cls.current_code_section = cls.current_code_section, section
+        yield
+        cls.current_code_section = prev_value
+
     def __init__(self, **kwargs):
         self._init_callstack = [
             frame_info for frame_info in inspect.stack()[1:]
             if Path(frame_info.filename).parent != CURRENT_DIR
         ]
         self._annotated_from = None
         self._comment = None
+        self._code_section = Synthesizable.current_code_section
 
     @property
     def name(self) -> str:
@@ -147,7 +172,6 @@ class Signal(Synthesizable):
     DEFAULT_DRIVER: str = "d"
 
     _new_signal_counter = count(0)
-    _signal_decl_in_verilog = False
     _str_with_net_name_only = False
 
     def __init__(
@@ -279,14 +303,6 @@ def with_width(self, width: int) -> Signal:
             return self[(-1,) * padding_size, :]
         return constant(0, padding_size, False) @ self
 
-    @classmethod
-    @contextmanager
-    def decl_in_verilog(cls):
-        """Declare a context to elaborate signals in Verilog style."""
-        prev_value, cls._signal_decl_in_verilog = cls._signal_decl_in_verilog, True
-        yield
-        cls._signal_decl_in_verilog = prev_value
-
     def signal_decl(self) -> str:
         """
         Declare the signal in the module implementation.
@@ -298,7 +314,14 @@ def signal_decl(self) -> str:
         if self.width == 0:
             raise ValueError("Signal width is not set and cannot be inferred")
 
-        template = SIGNAL_DECL_VERILOG_TEMPLATE if self._signal_decl_in_verilog else SIGNAL_DECL_TEMPLATE
+        match self._code_section:
+            case CodeSectionType.LOGIC:
+                template = SIGNAL_DECL_TEMPLATE
+            case CodeSectionType.VERILOG:
+                template = SIGNAL_DECL_VERILOG_TEMPLATE
+            case CodeSectionType.FORMAL:
+                template = SIGNAL_DECL_FORMAL_TEMPLATE
+
         decl = template.substitute(
             signed="signed" if self.signed else "",
             width=f"[{width - 1}:0]" if (width := self.width) > 1 else "",