Toward cleaning up variables and expressions in the composer

src/pystencilssfg/composer/basic_composer.py

 from __future__ import annotations
-from typing import Sequence
+from typing import Sequence, TypeAlias
 from abc import ABC, abstractmethod
 import numpy as np
 import sympy as sp
 from functools import reduce
 
-from pystencils import Field
+from pystencils import Field, TypedSymbol
 from pystencils.backend import KernelParameter, KernelFunction
-from pystencils.types import create_type, UserTypeSpec, PsCustomType, PsPointerType
+from pystencils.types import (
+    create_type,
+    UserTypeSpec,
+    PsCustomType,
+    PsPointerType,
+    PsType,
+)
 
 from ..context import SfgContext
 from .custom import CustomGenerator
@@ -58,8 +64,23 @@ class SfgNodeBuilder(ABC):
         pass
 
 
-ExprLike = str | SfgVar | AugExpr
-SequencerArg = tuple | str | AugExpr | SfgCallTreeNode | SfgNodeBuilder
+_ExprLike = (str, AugExpr, TypedSymbol)
+ExprLike: TypeAlias = str | AugExpr | TypedSymbol
+"""Things that may act as a C++ expression.
+
+Expressions need not necesserily have a known data type.
+"""
+
+_VarLike = (TypedSymbol, AugExpr)
+VarLike: TypeAlias = TypedSymbol | AugExpr
+"""Things that may act as a variable.
+
+Variables must always define their name *and* data type.
+"""
+
+_SequencerArg = (tuple, ExprLike, SfgCallTreeNode, SfgNodeBuilder)
+SequencerArg: TypeAlias = tuple | ExprLike | SfgCallTreeNode | SfgNodeBuilder
+"""Valid arguments to `make_sequence` and any sequencer that uses it."""
 
 
 class SfgBasicComposer(SfgIComposer):
@@ -208,7 +229,11 @@ class SfgBasicComposer(SfgIComposer):
         num_blocks_str = str(num_blocks)
         tpb_str = str(threads_per_block)
         stream_str = str(stream) if stream is not None else None
-        depends = _depends(num_blocks) | _depends(threads_per_block) | _depends(stream)
+
+        depends = _depends(num_blocks) | _depends(threads_per_block)
+        if stream is not None:
+            depends |= _depends(stream)
+
         return SfgCudaKernelInvocation(
             kernel_handle, num_blocks_str, tpb_str, stream_str, depends
         )
@@ -217,9 +242,9 @@ class SfgBasicComposer(SfgIComposer):
         """Syntax sequencing. For details, see `make_sequence`"""
         return make_sequence(*args)
 
-    def params(self, *args: SfgVar) -> SfgFunctionParams:
+    def params(self, *args: AugExpr) -> SfgFunctionParams:
         """Use inside a function body to add parameters to the function."""
-        return SfgFunctionParams(args)
+        return SfgFunctionParams([x.as_variable() for x in args])
 
     def require(self, *includes: str | SfgHeaderInclude) -> SfgRequireIncludes:
         return SfgRequireIncludes(
@@ -232,7 +257,7 @@ class SfgBasicComposer(SfgIComposer):
         ptr: bool = False,
         ref: bool = False,
         const: bool = False,
-    ):
+    ) -> PsType:
         if ptr and ref:
             raise SfgException("Create either a pointer, or a ref type, not both!")
 
@@ -250,11 +275,11 @@ class SfgBasicComposer(SfgIComposer):
         else:
             return base_type
 
-    def var(self, name: str, dtype: UserTypeSpec) -> SfgVar:
+    def var(self, name: str, dtype: UserTypeSpec) -> AugExpr:
         """Create a variable with given name and data type."""
-        return SfgVar(name, create_type(dtype))
+        return AugExpr(create_type(dtype)).var(name)
 
-    def init(self, lhs: SfgVar) -> SfgInplaceInitBuilder:
+    def init(self, lhs: VarLike) -> SfgInplaceInitBuilder:
         """Create a C++ in-place initialization.
 
         Usage:
@@ -270,7 +295,7 @@ class SfgBasicComposer(SfgIComposer):
 
             SomeClass obj { arg1, arg2, arg3 };
         """
-        return SfgInplaceInitBuilder(lhs)
+        return SfgInplaceInitBuilder(_asvar(lhs))
 
     def expr(self, fmt: str, *deps, **kwdeps):
         return AugExpr.format(fmt, *deps, **kwdeps)
@@ -329,15 +354,8 @@ class SfgBasicComposer(SfgIComposer):
 
 
 def make_statements(arg: ExprLike) -> SfgStatements:
-    match arg:
-        case str():
-            return SfgStatements(arg, (), ())
-        case SfgVar(name, _):
-            return SfgStatements(name, (), (arg,))
-        case AugExpr():
-            return SfgStatements(str(arg), (), arg.depends)
-        case _:
-            assert False
+    depends = _depends(arg)
+    return SfgStatements(str(arg), (), depends)
 
 
 def make_sequence(*args: SequencerArg) -> SfgSequence:
@@ -392,10 +410,8 @@ def make_sequence(*args: SequencerArg) -> SfgSequence:
             children.append(arg.resolve())
         elif isinstance(arg, SfgCallTreeNode):
             children.append(arg)
-        elif isinstance(arg, AugExpr):
-            children.append(SfgStatements(str(arg), (), arg.depends))
-        elif isinstance(arg, str):
-            children.append(SfgStatements(arg, (), ()))
+        elif isinstance(arg, _ExprLike):
+            children.append(make_statements(arg))
         elif isinstance(arg, tuple):
             #   Tuples are treated as blocks
             subseq = make_sequence(*arg)
@@ -407,25 +423,20 @@ def make_sequence(*args: SequencerArg) -> SfgSequence:
 
 
 class SfgInplaceInitBuilder(SfgNodeBuilder):
-    def __init__(self, lhs: SfgVar | AugExpr) -> None:
-        if isinstance(lhs, AugExpr):
-            lhs = lhs.as_variable()
-
+    def __init__(self, lhs: SfgVar) -> None:
         self._lhs: SfgVar = lhs
         self._depends: set[SfgVar] = set()
         self._rhs: str | None = None
 
     def __call__(
         self,
-        *rhs: str | AugExpr,
+        *rhs: ExprLike,
     ) -> SfgInplaceInitBuilder:
         if self._rhs is not None:
             raise SfgException("Assignment builder used multiple times.")
 
         self._rhs = ", ".join(str(expr) for expr in rhs)
-        self._depends = reduce(
-            set.union, (obj.depends for obj in rhs if isinstance(obj, AugExpr)), set()
-        )
+        self._depends = reduce(set.union, (_depends(obj) for obj in rhs), set())
         return self
 
     def resolve(self) -> SfgCallTreeNode:
@@ -538,12 +549,30 @@ def struct_from_numpy_dtype(
     return cls
 
 
-def _depends(expr: ExprLike | Sequence[ExprLike] | None) -> set[SfgVar]:
+def _asvar(var: VarLike) -> SfgVar:
+    match var:
+        case AugExpr():
+            return var.as_variable()
+        case TypedSymbol():
+            from pystencils import DynamicType
+
+            if isinstance(var.dtype, DynamicType):
+                raise SfgException(
+                    f"Unable to cast dynamically typed symbol {var} to a variable.\n"
+                    f"{var} has dynamic type {var.dtype}, which cannot be resolved to a type outside of a kernel."
+                )
+
+            return SfgVar(var.name, var.dtype)
+        case _:
+            raise ValueError(f"Invalid variable: {var}")
+
+
+def _depends(expr: ExprLike) -> set[SfgVar]:
     match expr:
         case None | str():
             return set()
-        case SfgVar():
-            return {expr}
+        case TypedSymbol():
+            return {_asvar(expr)}
         case AugExpr():
             return expr.depends
         case _:

src/pystencilssfg/composer/class_composer.py

 from __future__ import annotations
 from typing import Sequence
 
+from pystencils import TypedSymbol
 from pystencils.types import PsCustomType, UserTypeSpec
 
+from ..lang import AugExpr
 from ..ir import SfgCallTreeNode
 from ..ir.source_components import (
     SfgClass,
@@ -14,12 +16,18 @@ from ..ir.source_components import (
     SfgClassKeyword,
     SfgVisibility,
     SfgVisibilityBlock,
-    SfgVar,
 )
 from ..exceptions import SfgException
 
 from .mixin import SfgComposerMixIn
-from .basic_composer import SfgNodeBuilder, make_sequence
+from .basic_composer import (
+    SfgNodeBuilder,
+    make_sequence,
+    _VarLike,
+    VarLike,
+    ExprLike,
+    _asvar,
+)
 
 
 class SfgClassComposer(SfgComposerMixIn):
@@ -46,7 +54,7 @@ class SfgClassComposer(SfgComposerMixIn):
         def __call__(
             self,
             *args: (
-                SfgClassMember | SfgClassComposer.ConstructorBuilder | SfgVar | str
+                SfgClassMember | SfgClassComposer.ConstructorBuilder | VarLike | str
             ),
         ):
             for arg in args:
@@ -63,15 +71,21 @@ class SfgClassComposer(SfgComposerMixIn):
         Returned by `constructor`.
         """
 
-        def __init__(self, *params: SfgVar):
-            self._params = params
+        def __init__(self, *params: VarLike):
+            self._params = tuple(_asvar(p) for p in params)
             self._initializers: list[str] = []
             self._body: str | None = None
 
-        def init(self, initializer: str) -> SfgClassComposer.ConstructorBuilder:
+        def init(self, var: VarLike):
             """Add an initialization expression to the constructor's initializer list."""
-            self._initializers.append(initializer)
-            return self
+
+            def init_sequencer(expr: ExprLike):
+                expr = str(expr)
+                initializer = f"{_asvar(var)}{{ {expr} }}"
+                self._initializers.append(initializer)
+                return self
+
+            return init_sequencer
 
         def body(self, body: str):
             """Define the constructor body"""
@@ -120,7 +134,7 @@ class SfgClassComposer(SfgComposerMixIn):
         """Create a `private` visibility block in a class or struct body"""
         return SfgClassComposer.VisibilityContext(SfgVisibility.PRIVATE)
 
-    def constructor(self, *params: SfgVar):
+    def constructor(self, *params: VarLike):
         """In a class or struct body or visibility block, add a constructor.
 
         Args:
@@ -171,7 +185,7 @@ class SfgClassComposer(SfgComposerMixIn):
                 SfgClassComposer.VisibilityContext
                 | SfgClassMember
                 | SfgClassComposer.ConstructorBuilder
-                | SfgVar
+                | VarLike
                 | str
             ),
         ):
@@ -186,9 +200,9 @@ class SfgClassComposer(SfgComposerMixIn):
                     (
                         SfgClassMember,
                         SfgClassComposer.ConstructorBuilder,
-                        SfgVar,
                         str,
-                    ),
+                    )
+                    + _VarLike,
                 ):
                     if default_ended:
                         raise SfgException(
@@ -204,13 +218,17 @@ class SfgClassComposer(SfgComposerMixIn):
 
     @staticmethod
     def _resolve_member(
-        arg: SfgClassMember | SfgClassComposer.ConstructorBuilder | SfgVar | str,
-    ):
-        if isinstance(arg, SfgVar):
-            return SfgMemberVariable(arg.name, arg.dtype)
-        elif isinstance(arg, str):
-            return SfgInClassDefinition(arg)
-        elif isinstance(arg, SfgClassComposer.ConstructorBuilder):
-            return arg.resolve()
-        else:
-            return arg
+        arg: SfgClassMember | SfgClassComposer.ConstructorBuilder | VarLike | str,
+    ) -> SfgClassMember:
+        match arg:
+            case AugExpr() | TypedSymbol():
+                var = _asvar(arg)
+                return SfgMemberVariable(var.name, var.dtype)
+            case str():
+                return SfgInClassDefinition(arg)
+            case SfgClassComposer.ConstructorBuilder():
+                return arg.resolve()
+            case SfgClassMember():
+                return arg
+            case _:
+                raise ValueError(f"Invalid class member: {arg}")

src/pystencilssfg/lang/expressions.py

@@ -48,7 +48,7 @@ class DependentExpression:
 
     def __add__(self, other: DependentExpression):
         return DependentExpression(self.expr + other.expr, self.depends | other.depends)
-    
+
 
 class VarExpr(DependentExpression):
     def __init__(self, var: SfgVar):
@@ -61,6 +61,8 @@ class VarExpr(DependentExpression):
 
 
 class AugExpr:
+    __match_args__ = ("expr", "dtype")
+
     def __init__(self, dtype: PsType | None = None):
         self._dtype = dtype
         self._bound: DependentExpression | None = None
@@ -77,6 +79,7 @@ class AugExpr:
 
     @staticmethod
     def format(fmt: str, *deps, **kwdeps) -> AugExpr:
+        """Create a new `AugExpr` by combining existing expressions."""
         return AugExpr().bind(fmt, *deps, **kwdeps)
 
     def bind(self, fmt: str, *deps, **kwdeps):
@@ -109,11 +112,11 @@ class AugExpr:
             raise SfgException("This AugExpr has no known data type.")
 
         return self._dtype
-    
+
     @property
     def is_variable(self) -> bool:
         return isinstance(self._bound, VarExpr)
-    
+
     def as_variable(self) -> SfgVar:
         if not isinstance(self._bound, VarExpr):
             raise SfgException("This expression is not a variable")

tests/integration/expected/Variables.h

+#pragma once
+
+#include <cstdint>
+
+#define RESTRICT __restrict__
+
+class Scale {
+private:
+    float alpha;
+public:
+    Scale(float alpha) : alpha{ alpha } {}
+    void operator() (float *const _data_f, float *const _data_g);
+};

tests/integration/scripts/SimpleJacobi.py

@@ -20,4 +20,4 @@ with SourceFileGenerator() as sfg:
         sfg.map_field(u_dst, mdspan_ref(u_dst)),
         sfg.map_field(f, mdspan_ref(f)),
         sfg.call(poisson_kernel)
-    )
+    )
\ No newline at end of file

tests/integration/scripts/Variables.py

+import sympy as sp
+from pystencils import TypedSymbol, fields, kernel
+
+from pystencilssfg import SourceFileGenerator, SfgConfiguration
+
+with SourceFileGenerator() as sfg:
+    α = TypedSymbol("alpha", "float32")
+    f, g = fields("f, g: float32[10]")
+
+    @kernel
+    def scale():
+        f[0] @= α * g.center()
+
+    khandle = sfg.kernels.create(scale)
+
+    sfg.klass("Scale")(
+        sfg.private(α),
+        sfg.public(
+            sfg.constructor(α).init(α)(α.name),
+            sfg.method("operator()")(sfg.init(α)(f"this->{α}"), sfg.call(khandle)),
+        ),
+    )

tests/integration/test_generator_scripts.py

@@ -15,20 +15,45 @@ EXPECTED_DIR = path.join(THIS_DIR, "expected")
 @dataclass
 class ScriptInfo:
     script_name: str
+    """Name of the generator script, without .py-extension.
+    
+    Generator scripts must be located in the ``scripts`` folder.
+    """
+
     expected_outputs: tuple[str, ...]
+    """List of file extensions expected to be emitted by the generator script.
+    
+    Output files will all be placed in the ``out`` folder.
+    """
 
     compilable_output: str | None = None
+    """File extension of the output file that can be compiled.
+    
+    If this is set, and the expected file exists, the ``compile_cmd`` will be
+    executed to check for error-free compilation of the output.
+    """
+
     compile_cmd: str = f"g++ --std=c++17 -I {THIS_DIR}/deps/mdspan/include"
+    """Command to be invoked to compile the generated source file."""
 
 
 SCRIPTS = [
     ScriptInfo("SimpleJacobi", ("h", "cpp"), compilable_output="cpp"),
     ScriptInfo("SimpleClasses", ("h", "cpp")),
+    ScriptInfo("Variables", ("h", "cpp"), compilable_output="cpp"),
 ]
 
 
 @pytest.mark.parametrize("script_info", SCRIPTS)
 def test_generator_script(script_info: ScriptInfo):
+    """Test a generator script defined by ``script_info``.
+
+    The generator script will be run, with its output placed in the ``out`` folder.
+    If it is successful, its output files will be compared against
+    any files of the same name from the ``expected`` folder.
+    Finally, if any compilable files are specified, the test will attempt to compile them.
+    """
+
     script_name = script_info.script_name
     script_file = path.join(SCRIPTS_DIR, script_name + ".py")