hdl._ast: add SwitchValue, reimplement ArrayProxy with it.

Author: wanda-phi

amaranth/hdl/_ast.py

@@ -18,7 +18,7 @@
 __all__ = [
     "SyntaxError", "SyntaxWarning",
     "Shape", "signed", "unsigned", "ShapeCastable", "ShapeLike",
-    "Value", "Const", "C", "AnyConst", "AnySeq", "Operator", "Mux", "Part", "Slice", "Cat", "Concat",
+    "Value", "Const", "C", "AnyConst", "AnySeq", "Operator", "Mux", "Part", "Slice", "Cat", "Concat", "SwitchValue",
     "Array", "ArrayProxy",
     "Signal", "ClockSignal", "ResetSignal",
     "ValueCastable", "ValueLike",
@@ -1892,6 +1892,60 @@ def __repr__(self):
         return "(cat {})".format(" ".join(map(repr, self.parts)))
 
 
+@final
+class SwitchValue(Value):
+    def __init__(self, test, cases, *, src_loc=None, src_loc_at=0):
+        if src_loc is None:
+            super().__init__(src_loc_at=src_loc_at)
+        else:
+            self.src_loc = src_loc
+        self._test = Value.cast(test)
+        new_cases = []
+        for patterns, value in cases:
+            if patterns is not None:
+                if not isinstance(patterns, tuple):
+                    patterns = (patterns,)
+                new_patterns = ()
+                key_mask = (1 << len(self.test)) - 1
+                for key in _normalize_patterns(patterns, self._test.shape()):
+                    if isinstance(key, int):
+                        key = to_binary(key & key_mask, len(self.test))
+                    new_patterns = (*new_patterns, key)
+            else:
+                new_patterns = None
+            new_cases.append((new_patterns, Value.cast(value)))
+        self._cases = tuple(new_cases)
+
+    @property
+    def test(self):
+        return self._test
+
+    @property
+    def cases(self):
+        return self._cases
+
+    def shape(self):
+        return Shape._unify(value.shape() for _patterns, value in self._cases)
+
+    def _lhs_signals(self):
+        return union((value._lhs_signals() for _patterns, value in self.cases), start=SignalSet())
+
+    def _rhs_signals(self):
+        signals = union((value._rhs_signals() for _patterns, value in self.cases), start=SignalSet())
+        return self.test._rhs_signals() | signals
+
+    def __repr__(self):
+        def case_repr(patterns, value):
+            if patterns is None:
+                return f"(default {value!r})"
+            elif len(patterns) == 1:
+                return f"(case {patterns[0]} {value!r})"
+            else:
+                return "(case ({}) {!r})".format(" ".join(patterns), value)
+        case_reprs = (case_repr(patterns, value) for patterns, value in self.cases)
+        return "(switch-value {!r} {})".format(self.test, " ".join(case_reprs))
+
+
 class _SignalMeta(ABCMeta):
     def __call__(cls, shape=None, src_loc_at=0, **kwargs):
         signal = super().__call__(shape, **kwargs, src_loc_at=src_loc_at + 1)
@@ -2356,10 +2410,17 @@ def __repr__(self):
                                        ", ".join(map(repr, self._inner)))
 
 
+def _proxy_value(name):
+    @functools.wraps(getattr(Value, name))
+    def inner(self, *args, **kwargs):
+        return getattr(Value.cast(self), name)(*args, **kwargs)
+    return inner
+
+
 @final
-class ArrayProxy(Value):
+class ArrayProxy(ValueCastable):
     def __init__(self, elems, index, *, src_loc_at=0):
-        super().__init__(src_loc_at=1 + src_loc_at)
+        self.src_loc = tracer.get_src_loc(1 + src_loc_at)
         self._elems = elems
         self._index = Value.cast(index)
 
@@ -2385,19 +2446,73 @@ def shape(self):
         # elements. I.e., shape-wise, an array proxy must be identical to an equivalent mux tree.
         return Shape._unify(elem.shape() for elem in self._iter_as_values())
 
-    def _lhs_signals(self):
-        signals = union((elem._lhs_signals() for elem in self._iter_as_values()),
-                        start=SignalSet())
-        return signals
+    def as_value(self):
+        return SwitchValue(
+            self._index,
+            (
+                (index, value)
+                for index, value in enumerate(self._elems)
+                if index in range(1 << len(self._index))
+            ),
+            src_loc=self.src_loc,
+        )
 
-    def _rhs_signals(self):
-        signals = union((elem._rhs_signals() for elem in self._iter_as_values()),
-                        start=SignalSet())
-        return self.index._rhs_signals() | signals
+    def eq(self, value, *, src_loc_at=0):
+        return self.as_value().eq(value, src_loc_at=1 + src_loc_at)
 
     def __repr__(self):
         return "(proxy (array [{}]) {!r})".format(", ".join(map(repr, self.elems)), self.index)
 
+    as_signed = _proxy_value("as_signed")
+    as_unsigned = _proxy_value("as_unsigned")
+    __len__ = _proxy_value("__len__")
+    __bool__ = _proxy_value("__bool__")
+    bool = _proxy_value("bool")
+    __pos__ = _proxy_value("__pos__")
+    __neg__ = _proxy_value("__neg__")
+    __add__ = _proxy_value("__add__")
+    __radd__ = _proxy_value("__radd__")
+    __sub__ = _proxy_value("__sub__")
+    __rsub__ = _proxy_value("__rsub__")
+    __mul__ = _proxy_value("__mul__")
+    __rmul__ = _proxy_value("__rmul__")
+    __floordiv__ = _proxy_value("__floordiv__")
+    __rfloordiv__ = _proxy_value("__rfloordiv__")
+    __mod__ = _proxy_value("__mod__")
+    __rmod__ = _proxy_value("__rmod__")
+    __eq__ = _proxy_value("__eq__")
+    __ne__ = _proxy_value("__ne__")
+    __lt__ = _proxy_value("__lt__")
+    __le__ = _proxy_value("__le__")
+    __gt__ = _proxy_value("__gt__")
+    __ge__ = _proxy_value("__ge__")
+    __abs__ = _proxy_value("__abs__")
+    __invert__ = _proxy_value("__invert__")
+    __and__ = _proxy_value("__and__")
+    __rand__ = _proxy_value("__rand__")
+    __or__ = _proxy_value("__or__")
+    __ror__ = _proxy_value("__ror__")
+    __xor__ = _proxy_value("__xor__")
+    __rxor__ = _proxy_value("__rxor__")
+    any = _proxy_value("any")
+    all = _proxy_value("all")
+    xor = _proxy_value("xor")
+    implies = _proxy_value("implies")
+    __lshift__ = _proxy_value("__lshift__")
+    __rlshift__ = _proxy_value("__rlshift__")
+    __rshift__ = _proxy_value("__rshift__")
+    __rrshift__ = _proxy_value("__rrshift__")
+    shift_left = _proxy_value("shift_left")
+    shift_right = _proxy_value("shift_right")
+    rotate_left = _proxy_value("rotate_left")
+    rotate_right = _proxy_value("rotate_right")
+    __contains__ = _proxy_value("__contains__")
+    bit_select = _proxy_value("bit_select")
+    word_select = _proxy_value("word_select")
+    replicate = _proxy_value("replicate")
+    matches = _proxy_value("matches")
+    __format__ = _proxy_value("__format__")
+
 
 @final
 class Initial(Value):
@@ -2772,7 +2887,7 @@ def __init__(self, test, cases, *, src_loc=None, src_loc_at=0):
             self.src_loc = src_loc
 
         self._test  = Value.cast(test)
-        self._cases = []
+        new_cases = []
         for patterns, stmts, case_src_loc in cases:
             if patterns is not None:
                 # Map: key -> (key,); (key...) -> (key...)
@@ -2787,10 +2902,8 @@ def __init__(self, test, cases, *, src_loc=None, src_loc_at=0):
                     new_patterns = (*new_patterns, key)
             else:
                 new_patterns = None
-            if not isinstance(stmts, Iterable):
-                stmts = [stmts]
-            self._cases.append((new_patterns, Statement.cast(stmts), case_src_loc))
-        self._cases = tuple(self._cases)
+            new_cases.append((new_patterns, Statement.cast(stmts), case_src_loc))
+        self._cases = tuple(new_cases)
 
     @property
     def test(self):
@@ -2816,7 +2929,7 @@ def case_repr(patterns, stmts):
                 return f"(case {patterns[0]} {stmts_repr})"
             else:
                 return "(case ({}) {})".format(" ".join(patterns), stmts_repr)
-        case_reprs = [case_repr(patterns, stmts) for patterns, stmts, _src_loc in self.cases]
+        case_reprs = (case_repr(patterns, stmts) for patterns, stmts, _src_loc in self.cases)
         return "(switch {!r} {})".format(self.test, " ".join(case_reprs))
 
 

amaranth/hdl/_ir.py

@@ -885,30 +885,39 @@ def emit_rhs(self, module_idx: int, value: _ast.Value) -> Tuple[_nir.Value, bool
                              stride=value.stride, offset=offset, src_loc=value.src_loc)
             result = self.netlist.add_value_cell(value.width, cell)
             signed = False
-        elif isinstance(value, _ast.ArrayProxy):
-            elems = [self.emit_rhs(module_idx, elem) for elem in value.elems]
+        elif isinstance(value, _ast.SwitchValue):
+            test, _signed = self.emit_rhs(module_idx, value.test)
+            conds = []
+            elems = []
+            for patterns, elem, in value.cases:
+                if patterns is not None:
+                    if not patterns:
+                        # Hack: empty pattern set cannot be supported by RTLIL.
+                        continue
+                    for pattern in patterns:
+                        assert len(pattern) == len(test)
+                    cell = _nir.Matches(module_idx, value=test, patterns=patterns,
+                                        src_loc=value.src_loc)
+                    net, = self.netlist.add_value_cell(1, cell)
+                    conds.append(net)
+                else:
+                    conds.append(_nir.Net.from_const(1))
+                elems.append(self.emit_rhs(module_idx, elem))
+            cell = _nir.PriorityMatch(module_idx, en=_nir.Net.from_const(1),
+                                      inputs=_nir.Value(conds),
+                                      src_loc=value.src_loc)
+            conds = self.netlist.add_value_cell(len(conds), cell)
             shape = _ast.Shape._unify(
                 _ast.Shape(len(value), signed)
                 for value, signed in elems
             )
             elems = tuple(self.extend(elem, elem_signed, shape.width) for elem, elem_signed in elems)
-            index, _signed = self.emit_rhs(module_idx, value.index)
-            conds = []
-            for case_index in range(len(elems)):
-                cell = _nir.Matches(module_idx, value=index,
-                                       patterns=(to_binary(case_index, len(index)),),
-                                       src_loc=value.src_loc)
-                subcond, = self.netlist.add_value_cell(1, cell)
-                conds.append(subcond)
-            conds = _nir.Value(conds)
-            cell = _nir.PriorityMatch(module_idx, en=_nir.Net.from_const(1), inputs=conds, src_loc=value.src_loc)
-            conds = self.netlist.add_value_cell(len(conds), cell)
             assignments = [
-                _nir.Assignment(cond=cond, start=0, value=elem, src_loc=value.src_loc)
-                for cond, elem in zip(conds, elems)
+                _nir.Assignment(cond=subcond, start=0, value=elem, src_loc=value.src_loc)
+                for subcond, elem in zip(conds, elems)
             ]
-            cell = _nir.AssignmentList(module_idx, default=elems[0], assignments=assignments,
-                                 src_loc=value.src_loc)
+            cell = _nir.AssignmentList(module_idx, default=_nir.Value.from_const(0, shape.width),
+                                       assignments=assignments, src_loc=value.src_loc)
             result = self.netlist.add_value_cell(shape.width, cell)
             signed = shape.signed
         elif isinstance(value, _ast.Concat):
@@ -1017,19 +1026,29 @@ def emit_assign(self, module_idx: int, cd: "_cd.ClockDomain | None", lhs: _ast.V
                 else:
                     subrhs = rhs
                 self.emit_assign(module_idx, cd, lhs.value, start, subrhs, subcond, src_loc=src_loc)
-        elif isinstance(lhs, _ast.ArrayProxy):
-            index, _signed = self.emit_rhs(module_idx, lhs.index)
+        elif isinstance(lhs, _ast.SwitchValue):
+            test, _signed = self.emit_rhs(module_idx, lhs.test)
             conds = []
-            for case_index in range(len(lhs.elems)):
-                cell = _nir.Matches(module_idx, value=index,
-                                       patterns=(to_binary(case_index, len(index)),),
-                                       src_loc=lhs.src_loc)
-                subcond, = self.netlist.add_value_cell(1, cell)
-                conds.append(subcond)
+            elems = []
+            for patterns, elem in lhs.cases:
+                if patterns is not None:
+                    if not patterns:
+                        # Hack: empty pattern set cannot be supported by RTLIL.
+                        continue
+                    for pattern in patterns:
+                        assert len(pattern) == len(test)
+                    cell = _nir.Matches(module_idx, value=test, patterns=patterns,
+                                        src_loc=lhs.src_loc)
+                    net, = self.netlist.add_value_cell(1, cell)
+                    conds.append(net)
+                else:
+                    conds.append(_nir.Net.from_const(1))
+                elems.append(elem)
             conds = _nir.Value(conds)
-            cell = _nir.PriorityMatch(module_idx, en=cond, inputs=conds, src_loc=lhs.src_loc)
+            cell = _nir.PriorityMatch(module_idx, en=_nir.Net.from_const(1),
+                                      inputs=conds, src_loc=lhs.src_loc)
             conds = self.netlist.add_value_cell(len(conds), cell)
-            for subcond, val in zip(conds, lhs.elems):
+            for subcond, val in zip(conds, elems):
                 self.emit_assign(module_idx, cd, val, lhs_start, rhs[:len(val)], subcond, src_loc=src_loc)
         elif isinstance(lhs, _ast.Operator):
             assert lhs.operator in ('u', 's')

amaranth/hdl/_xfrm.py

@@ -58,7 +58,7 @@ def on_Concat(self, value):
         pass # :nocov:
 
     @abstractmethod
-    def on_ArrayProxy(self, value):
+    def on_SwitchValue(self, value):
         pass # :nocov:
 
     @abstractmethod
@@ -90,8 +90,8 @@ def on_value(self, value):
             new_value = self.on_Part(value)
         elif type(value) is Concat:
             new_value = self.on_Concat(value)
-        elif type(value) is ArrayProxy:
-            new_value = self.on_ArrayProxy(value)
+        elif type(value) is SwitchValue:
+            new_value = self.on_SwitchValue(value)
         elif type(value) is Initial:
             new_value = self.on_Initial(value)
         else:
@@ -133,9 +133,8 @@ def on_Part(self, value):
     def on_Concat(self, value):
         return Concat(self.on_value(o) for o in value.parts)
 
-    def on_ArrayProxy(self, value):
-        return ArrayProxy([self.on_value(elem) for elem in value._iter_as_values()],
-                          self.on_value(value.index))
+    def on_SwitchValue(self, value):
+        return SwitchValue(self.on_value(value.test), [(patterns, self.on_value(val)) for patterns, val in value.cases])
 
     def on_Initial(self, value):
         return value
@@ -399,10 +398,10 @@ def on_Concat(self, value):
         for o in value.parts:
             self.on_value(o)
 
-    def on_ArrayProxy(self, value):
-        for elem in value._iter_as_values():
-            self.on_value(elem)
-        self.on_value(value.index)
+    def on_SwitchValue(self, value):
+        self.on_value(value.test)
+        for patterns, val in value.cases:
+            self.on_value(val)
 
     def on_Initial(self, value):
         pass