From ded84fe9d640c1fffa4a089fcb7233319da5ea93 Mon Sep 17 00:00:00 2001
From: Jaro Habiger <jarohabiger@googlemail.com>
Date: Fri, 5 Jan 2024 14:57:27 +0100
Subject: [PATCH] sim: fix ValueCastable not being recognized as a coroutine
 command

---
 amaranth/sim/_pycoro.py |  4 +++-
 tests/test_sim.py       | 21 +++++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/amaranth/sim/_pycoro.py b/amaranth/sim/_pycoro.py
index 10ce1ca..e2714a2 100644
--- a/amaranth/sim/_pycoro.py
+++ b/amaranth/sim/_pycoro.py
@@ -1,7 +1,7 @@
 import inspect
 
 from ..hdl import *
-from ..hdl.ast import Statement, SignalSet
+from ..hdl.ast import Statement, SignalSet, ValueCastable
 from .core import Tick, Settle, Delay, Passive, Active
 from ._base import BaseProcess
 from ._pyrtl import _ValueCompiler, _RHSValueCompiler, _StatementCompiler
@@ -66,6 +66,8 @@ class PyCoroProcess(BaseProcess):
                     command = self.default_cmd
                 response = None
 
+                if isinstance(command, ValueCastable):
+                    command = Value.cast(command)
                 if isinstance(command, Value):
                     exec(_RHSValueCompiler.compile(self.state, command, mode="curr"),
                         self.exec_locals)
diff --git a/tests/test_sim.py b/tests/test_sim.py
index 19e493b..6f561a4 100644
--- a/tests/test_sim.py
+++ b/tests/test_sim.py
@@ -690,6 +690,27 @@ class SimulatorIntegrationTestCase(FHDLTestCase):
             sim.add_process(process)
         self.assertTrue(survived)
 
+    def test_value_castable(self):
+        class MyValue(ValueCastable):
+            @ValueCastable.lowermethod
+            def as_value(self):
+                return Signal()
+            
+            def shape():
+                return unsigned(1)
+        
+        a = Array([1,2,3])
+        a[MyValue()]
+
+        survived = False
+        with self.assertSimulation(Module()) as sim:
+            def process():
+                nonlocal survived
+                yield MyValue()
+                survived = True
+            sim.add_process(process)
+        self.assertTrue(survived)
+
     def setUp_memory(self, rd_synchronous=True, rd_transparent=True, wr_granularity=None):
         self.m = Module()
         self.memory = Memory(width=8, depth=4, init=[0xaa, 0x55])