split namespaces into an entity and an element part

src/pystencilssfg/ir/source_components.py

@@ -39,9 +39,9 @@ class SfgCodeEntity:
     Each code entity has a name and an optional enclosing namespace.
     """
 
-    def __init__(self, name: str, namespace: SfgNamespace | None) -> None:
+    def __init__(self, name: str, parent_namespace: SfgNamespace | None) -> None:
         self._name = name
-        self._namespace: SfgNamespace | None = namespace
+        self._namespace: SfgNamespace | None = parent_namespace
 
     @property
     def name(self) -> str:
@@ -57,11 +57,23 @@ class SfgCodeEntity:
             return self._name
 
     @property
-    def namespace(self) -> SfgNamespace | None:
+    def parent_namespace(self) -> SfgNamespace | None:
         """Parent namespace of this entity"""
         return self._namespace
 
 
+class SfgNamespace(SfgCodeEntity):
+    """A C++ namespace.
+
+    Each namespace has a `name` and a `parent`; its fully qualified name is given as
+    ``<parent.name>::<name>``.
+
+    Args:
+        name: Local name of this namespace
+        parent: Parent namespace enclosing this namespace
+    """
+
+
 class SfgKernelHandle(SfgCodeEntity):
     """Handle to a pystencils kernel."""
 
@@ -100,7 +112,76 @@ class SfgKernelHandle(SfgCodeEntity):
         return self._kernel
 
 
+class SfgKernelNamespace(SfgNamespace):
+    """A namespace grouping together a number of kernels."""
+
+    def __init__(self, name: str, parent: SfgNamespace | None):
+        super().__init__(name, parent)
+        self._kernels: dict[str, SfgKernelHandle] = []
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def kernels(self) -> tuple[SfgKernelHandle, ...]:
+        return tuple(self._kernels.values())
+
+    def add(self, kernel: Kernel, name: str | None = None):
+        """Adds an existing pystencils AST to this namespace.
+        If a name is specified, the AST's function name is changed."""
+        if name is None:
+            kernel_name = kernel.name
+        else:
+            kernel_name = name
+
+        if kernel_name in self._kernels:
+            raise ValueError(
+                f"Duplicate kernels: A kernel called {kernel_name} already exists in namespace {self.fqname}"
+            )
+
+        if name is not None:
+            kernel.name = kernel_name
+
+        khandle = SfgKernelHandle(kernel_name, self, kernel)
+        self._kernels[kernel_name] = khandle
+
+        #   TODO: collect includes later
+        # for header in kernel.required_headers:
+        #     self._ctx.add_include(
+        #         SfgHeaderInclude(HeaderFile.parse(header), private=True)
+        #     )
+
+        return khandle
+
+    def create(
+        self,
+        assignments,
+        name: str | None = None,
+        config: CreateKernelConfig | None = None,
+    ):
+        """Creates a new pystencils kernel from a list of assignments and a configuration.
+        This is a wrapper around `pystencils.create_kernel`
+        with a subsequent call to `add`.
+        """
+        if config is None:
+            config = CreateKernelConfig()
+
+        if name is not None:
+            if name in self._kernels:
+                raise ValueError(
+                    f"Duplicate kernels: A kernel with name {name} already exists in namespace {self.fqname}"
+                )
+            config = replace(config, function_name=name)
+
+        # type: ignore
+        kernel = create_kernel(assignments, config=config)
+        return self.add(kernel)
+
+
 class SfgFunction(SfgCodeEntity):
+    """A free function."""
+
     __match_args__ = ("name", "tree", "parameters")
 
     def __init__(
@@ -142,6 +223,8 @@ class SfgFunction(SfgCodeEntity):
 
 
 class SfgVisibility(Enum):
+    """Visibility qualifiers of C++"""
+
     DEFAULT = auto()
     PRIVATE = auto()
     PROTECTED = auto()
@@ -160,6 +243,8 @@ class SfgVisibility(Enum):
 
 
 class SfgClassKeyword(Enum):
+    """Class keywords of C++"""
+
     STRUCT = auto()
     CLASS = auto()
 
@@ -172,6 +257,8 @@ class SfgClassKeyword(Enum):
 
 
 class SfgClassMember(ABC):
+    """Base class for class member entities"""
+
     def __init__(self) -> None:
         self._cls: SfgClass | None = None
         self._visibility: SfgVisibility | None = None
@@ -205,25 +292,56 @@ class SfgClassMember(ABC):
 
 
 class SfgMemberVariable(SfgVar, SfgClassMember):
+    """Variable that is a field of a class"""
+
     def __init__(self, name: str, dtype: PsType):
         SfgVar.__init__(self, name, dtype)
         SfgClassMember.__init__(self)
 
 
-class SfgMethod(SfgFunction, SfgClassMember):
+class SfgMethod(SfgClassMember):
+    """Instance method of a class"""
+
+    __match_args__ = ("name", "tree", "parameters")
+
     def __init__(
         self,
         name: str,
         tree: SfgCallTreeNode,
-        return_type: PsType = PsCustomType("void"),
+        return_type: PsType = void,
         inline: bool = False,
         const: bool = False,
     ):
-        SfgFunction.__init__(self, name, tree, return_type=return_type, inline=inline)
-        SfgClassMember.__init__(self)
+        super().__init__()
 
+        self._name = name
+        self._tree = tree
+        self._return_type = return_type
+        self._inline = inline
         self._const = const
-        self._parameters: set[SfgVar] = set()
+
+        self._parameters: set[SfgVar]
+
+        from .postprocessing import CallTreePostProcessing
+
+        param_collector = CallTreePostProcessing()
+        self._parameters = param_collector(self._tree).function_params
+
+    @property
+    def parameters(self) -> set[SfgVar]:
+        return self._parameters
+
+    @property
+    def tree(self) -> SfgCallTreeNode:
+        return self._tree
+
+    @property
+    def return_type(self) -> PsType:
+        return self._return_type
+
+    @property
+    def inline(self) -> bool:
+        return self._inline
 
     @property
     def const(self) -> bool:
@@ -231,6 +349,8 @@ class SfgMethod(SfgFunction, SfgClassMember):
 
 
 class SfgConstructor(SfgClassMember):
+    """Constructor of a class"""
+
     __match_args__ = ("parameters", "initializers", "body")
 
     def __init__(
@@ -503,7 +623,7 @@ class SfgVisibilityBlock:
         self._cls = cls
 
 
-class SfgNamespace:
+class SfgNamespaceDef:
     """A C++ namespace.
 
     Each namespace has a `name` and a `parent`; its fully qualified name is given as
@@ -514,24 +634,10 @@ class SfgNamespace:
         parent: Parent namespace enclosing this namespace
     """
 
-    def __init__(self, name: str, parent: SfgNamespace | None) -> None:
-        self._name: str = name
-        self._parent: SfgNamespace | None = parent
+    def __init__(self, namespace: SfgNamespace) -> None:
+        self._namespace = namespace
         self._elements: list[SfgNamespaceElement] = []
 
-    @property
-    def name(self) -> str:
-        """The name of this namespace"""
-        return self._name
-
-    @property
-    def fqname(self) -> str:
-        """The fully qualified name of this namespace"""
-        if self._parent is not None:
-            return self._parent.fqname + "::" + self._name
-        else:
-            return self._name
-
     @property
     def elements(self) -> list[SfgNamespaceElement]:
         """Sequence of source elements that make up the body of this namespace"""
@@ -542,74 +648,7 @@ class SfgNamespace:
         self._elements = list(elems)
 
 
-class SfgKernelNamespace(SfgNamespace):
-    """A namespace grouping together a number of kernels."""
-
-    def __init__(self, name: str, parent: SfgNamespace | None):
-        super().__init__(name, parent)
-        self._kernels: dict[str, SfgKernelHandle] = []
-
-    @property
-    def name(self):
-        return self._name
-
-    @property
-    def kernels(self) -> tuple[SfgKernelHandle, ...]:
-        return tuple(self._kernels.values())
-
-    def add(self, kernel: Kernel, name: str | None = None):
-        """Adds an existing pystencils AST to this namespace.
-        If a name is specified, the AST's function name is changed."""
-        if name is None:
-            kernel_name = kernel.name
-        else:
-            kernel_name = name
-
-        if kernel_name in self._kernels:
-            raise ValueError(
-                f"Duplicate kernels: A kernel called {kernel_name} already exists in namespace {self.fqname}"
-            )
-
-        if name is not None:
-            kernel.name = kernel_name
-
-        khandle = SfgKernelHandle(kernel_name, self, kernel)
-        self._kernels[kernel_name] = khandle
-
-        #   TODO: collect includes later
-        # for header in kernel.required_headers:
-        #     self._ctx.add_include(
-        #         SfgHeaderInclude(HeaderFile.parse(header), private=True)
-        #     )
-
-        return khandle
-
-    def create(
-        self,
-        assignments,
-        name: str | None = None,
-        config: CreateKernelConfig | None = None,
-    ):
-        """Creates a new pystencils kernel from a list of assignments and a configuration.
-        This is a wrapper around `pystencils.create_kernel`
-        with a subsequent call to `add`.
-        """
-        if config is None:
-            config = CreateKernelConfig()
-
-        if name is not None:
-            if name in self._kernels:
-                raise ValueError(
-                    f"Duplicate kernels: A kernel with name {name} already exists in namespace {self.fqname}"
-                )
-            config = replace(config, function_name=name)
-
-        # type: ignore
-        kernel = create_kernel(assignments, config=config)
-        return self.add(kernel)
-
-
-SfgNamespaceElement = str | SfgNamespace | SfgEntityDecl | SfgEntityDef
+SfgNamespaceElement = str | SfgNamespaceDef | SfgEntityDecl | SfgEntityDef
 """Elements that may be placed inside a namespace, including the global namespace."""