core: Make SSAValue generic

Adds a generic parameter to SSAValue that constrains the type associated
with the SSAValue.
This is useful whenever a function only accepts SSAValue of a given type.

stack-info: PR: #3987, branch: math-fehr/stack/2

tests/test_ssa_value.py

@@ -1,8 +1,9 @@
 import pytest
 
-from xdsl.dialects.builtin import StringAttr, i32
+from xdsl.dialects.builtin import IntegerType, StringAttr, i32
 from xdsl.ir import Block, BlockArgument, SSAValue
 from xdsl.irdl import IRDLOperation, irdl_op_definition, result_def
+from xdsl.utils.test_value import TestSSAValue
 
 
 @irdl_op_definition
@@ -59,3 +60,16 @@ def test_invalid_ssa_vals(name: str):
     val = BlockArgument(i32, Block(), 0)
     with pytest.raises(ValueError):
         val.name_hint = name
+
+
+def test_typed_ssavalue():
+    """
+    Check that SSAValue type hinting works as expected.
+    This is a type hinting test.
+    """
+    value = TestSSAValue(i32)
+
+    def foo(arg: SSAValue[IntegerType]):
+        pass
+
+    foo(value)

xdsl/ir/core.py

@@ -478,13 +478,13 @@ def remove_use(self, use: Use):
 
 
 @dataclass(eq=False)
-class SSAValue(IRWithUses, ABC):
+class SSAValue(Generic[AttributeCovT], IRWithUses, ABC):
     """
     A reference to an SSA variable.
     An SSA variable is either an operation result, or a basic block argument.
     """
 
-    type: Attribute
+    type: AttributeCovT
     """Each SSA variable is associated to a type."""
 
     _name: str | None = field(init=False, default=None)

xdsl/printer.py

@@ -119,6 +119,7 @@ def print(self, *argv: Any) -> None:
                 self.print_string(arg)
                 continue
             if isinstance(arg, SSAValue):
+                arg = cast(SSAValue[Attribute], arg)
                 self.print_ssa_value(arg)
                 continue
             if isinstance(arg, Attribute):

xdsl/utils/test_value.py

@@ -1,9 +1,11 @@
+from typing import Generic
+
 import pytest
 
-from xdsl.ir import Block, Operation, SSAValue
+from xdsl.ir import AttributeCovT, Block, Operation, SSAValue
 
 
-class TestSSAValue(SSAValue):
+class TestSSAValue(Generic[AttributeCovT], SSAValue[AttributeCovT]):
     @property
     def owner(self) -> Operation | Block:
         pytest.fail("Attempting to get the owner of a `TestSSAValue`")