From c4bbcc6f8a900fa74afcaa0b4ed11ea40f61bdce Mon Sep 17 00:00:00 2001
From: Wanda <wanda@phinode.net>
Date: Tue, 2 Apr 2024 12:04:51 +0200
Subject: [PATCH] hdl._ast: enforce the `ShapeCastable.const` contract in
 `Const()`.

---
 amaranth/hdl/_ast.py   | 8 +++++++-
 amaranth/lib/data.py   | 7 +------
 amaranth/lib/memory.py | 4 ++--
 amaranth/lib/wiring.py | 2 +-
 tests/test_hdl_ast.py  | 8 +++++++-
 5 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/amaranth/hdl/_ast.py b/amaranth/hdl/_ast.py
index 3b71bc1..15188c8 100644
--- a/amaranth/hdl/_ast.py
+++ b/amaranth/hdl/_ast.py
@@ -1503,7 +1503,13 @@ class ValueLike(metaclass=_ValueLikeMeta):
 class _ConstMeta(ABCMeta):
     def __call__(cls, value, shape=None, src_loc_at=0, **kwargs):
         if isinstance(shape, ShapeCastable):
-            return shape.const(value)
+            value = shape.const(value)
+            cast_shape = Shape.cast(shape)
+            cast_value = Const.cast(value)
+            if cast_value.shape() != cast_shape:
+                raise ValueError(f"Constant returned by {shape!r}.const() must have the shape that "
+                                 f"it casts to, {cast_shape!r}, and not {cast_value.shape()!r}")
+            return value
         return super().__call__(value, shape, **kwargs, src_loc_at=src_loc_at + 1)
 
 
diff --git a/amaranth/lib/data.py b/amaranth/lib/data.py
index a141ee2..5f9fe37 100644
--- a/amaranth/lib/data.py
+++ b/amaranth/lib/data.py
@@ -223,12 +223,7 @@ class Layout(ShapeCastable, metaclass=ABCMeta):
             field = self[key]
             cast_field_shape = Shape.cast(field.shape)
             if isinstance(field.shape, ShapeCastable):
-                key_value = hdl.Const.cast(field.shape.const(key_value))
-                if key_value.shape() != cast_field_shape:
-                    raise ValueError("Constant returned by {!r}.const() must have the shape that "
-                                     "it casts to, {!r}, and not {!r}"
-                                     .format(field.shape, cast_field_shape,
-                                             key_value.shape()))
+                key_value = hdl.Const.cast(hdl.Const(key_value, field.shape))
             elif not isinstance(key_value, hdl.Const):
                 key_value = hdl.Const(key_value, cast_field_shape)
             mask = ((1 << cast_field_shape.width) - 1) << field.offset
diff --git a/amaranth/lib/memory.py b/amaranth/lib/memory.py
index 2361ea9..ea3c3ce 100644
--- a/amaranth/lib/memory.py
+++ b/amaranth/lib/memory.py
@@ -83,7 +83,7 @@ class Memory(wiring.Component):
 
             if isinstance(shape, ShapeCastable):
                 self._elems = [None] * depth
-                self._raw = [Const.cast(shape.const(None)).value] * depth
+                self._raw = [Const.cast(Const(None, shape)).value] * depth
             else:
                 self._elems = [0] * depth
                 self._raw = self._elems # intentionally mutably aliased
@@ -113,7 +113,7 @@ class Memory(wiring.Component):
                     self[actual_index] = actual_value
             else:
                 if isinstance(self._shape, ShapeCastable):
-                    self._raw[index] = Const.cast(self._shape.const(value)).value
+                    self._raw[index] = Const.cast(Const(value, self._shape)).value
                 else:
                     value = operator.index(value)
                     # self._raw[index] assigned by the following line
diff --git a/amaranth/lib/wiring.py b/amaranth/lib/wiring.py
index 29c85e2..2fcad85 100644
--- a/amaranth/lib/wiring.py
+++ b/amaranth/lib/wiring.py
@@ -139,7 +139,7 @@ class Member:
             # TODO: We need a simpler way to check for "is this a valid constant initializer"
             if issubclass(type(self._description), ShapeCastable):
                 try:
-                    self._init_as_const = Const.cast(self._description.const(self._init))
+                    self._init_as_const = Const.cast(Const(self._init, self._description))
                 except Exception as e:
                     raise TypeError(f"Port member initial value {self._init!r} is not a valid "
                                     f"constant initializer for {self._description}") from e
diff --git a/tests/test_hdl_ast.py b/tests/test_hdl_ast.py
index 4c80e78..d3a20fa 100644
--- a/tests/test_hdl_ast.py
+++ b/tests/test_hdl_ast.py
@@ -522,10 +522,16 @@ class ConstTestCase(FHDLTestCase):
             hash(Const(0))
 
     def test_shape_castable(self):
-        class MockConstValue:
+        class MockConstValue(ValueCastable):
             def __init__(self, value):
                 self.value = value
 
+            def shape(self):
+                return MockConstShape()
+            
+            def as_value(self):
+                return Const(self.value, 8)
+
         class MockConstShape(ShapeCastable):
             def as_shape(self):
                 return unsigned(8)