From 2c98ad94b5af9887ce0d999ee331a21381e6d18f Mon Sep 17 00:00:00 2001
From: Florent Kermarrec <florent@enjoy-digital.fr>
Date: Fri, 15 Oct 2021 13:33:09 +0200
Subject: [PATCH] fhdl/verilog: Create _print_operator/_print_slice, move code
 outside _print_expression and cleanup/simplify.

---
 litex/gen/fhdl/verilog.py | 115 +++++++++++++++++++++++---------------
 1 file changed, 69 insertions(+), 46 deletions(-)

diff --git a/litex/gen/fhdl/verilog.py b/litex/gen/fhdl/verilog.py
index 7e98fe316..db0d8adf9 100644
--- a/litex/gen/fhdl/verilog.py
+++ b/litex/gen/fhdl/verilog.py
@@ -100,6 +100,73 @@ def _print_constant(node):
         value = abs(node.value),
     ), node.signed
 
+# Print Operator -----------------------------------------------------------------------------------
+
+(UNARY, BINARY, TERNARY) = (1, 2, 3)
+
+def _print_operator(ns, node):
+    operator = node.op
+    operands = node.operands
+    arity    = len(operands)
+    assert arity in [UNARY, BINARY, TERNARY]
+
+    def to_signed(r):
+        return f"$signed({{1'd0 {r}}}))"
+
+    # Unary Operator.
+    if arity == UNARY:
+        r1, s1 = _print_expression(ns, operands[0])
+        # Negation Operator.
+        if operator == "-":
+            # Negate and convert to signed if not already.
+            r = "-" + (r1 if s1 else to_signed(r1))
+            s = True
+        # Other Operators.
+        else:
+            r = operator + r1
+            s = s1
+
+    # Binary Operator.
+    if arity == BINARY:
+        r1, s1 = _print_expression(ns, operands[0])
+        r2, s2 = _print_expression(ns, operands[1])
+        # Convert all expressions to signed when at least one is signed.
+        if operator not in ["<<<", ">>>"]:
+            if s2 and not s1:
+                r1 = to_signed(r1)
+            if s1 and not s2:
+                r2 = to_signed(r2)
+        r = f"{r1} {operator} {r2}"
+        s = s1 or s2
+
+    # Ternary Operator.
+    if arity == TERNARY:
+        assert operator == "m"
+        r1, s1 = _print_expression(ns, operands[0])
+        r2, s2 = _print_expression(ns, operands[1])
+        r3, s3 = _print_expression(ns, operands[2])
+        # Convert all expressions to signed when at least one is signed.
+        if s2 and not s3:
+            r3 = to_signed(r3)
+        if s3 and not s2:
+            r2 = to_signed(r2)
+        r = f"{r1} ? {r2} : {r3}"
+        s = s2 or s3
+
+    return f"({r})", s
+
+# Print Slice --------------------------------------------------------------------------------------
+
+def _print_slice(ns, node):
+    assert (node.stop - node.start) >= 1
+    if (isinstance(node.value, Signal) and len(node.value) == 1):
+        assert node.start == 0
+        sr = "" # Avoid slicing 1-bit Signals.
+    else:
+        sr = f"[{node.stop-1}:{node.start}]" if (node.stop - node.start) > 1 else f"[{node.start}]"
+    r, s = _print_expression(ns, node.value)
+    return r + sr, s
+
 # Print Expression ---------------------------------------------------------------------------------
 
 def _print_expression(ns, node):
@@ -113,55 +180,11 @@ def _print_expression(ns, node):
 
     # Operator.
     elif isinstance(node, _Operator):
-        arity = len(node.operands)
-        r1, s1 = _print_expression(ns, node.operands[0])
-        if arity == 1:
-            if node.op == "-":
-                if s1:
-                    r = node.op + r1
-                else:
-                    r = "-$signed({1'd0, " + r1 + "})"
-                s = True
-            else:
-                r = node.op + r1
-                s = s1
-        elif arity == 2:
-            r2, s2 = _print_expression(ns, node.operands[1])
-            if node.op not in ["<<<", ">>>"]:
-                if s2 and not s1:
-                    r1 = "$signed({1'd0, " + r1 + "})"
-                if s1 and not s2:
-                    r2 = "$signed({1'd0, " + r2 + "})"
-            r = r1 + " " + node.op + " " + r2
-            s = s1 or s2
-        elif arity == 3:
-            assert node.op == "m"
-            r2, s2 = _print_expression(ns, node.operands[1])
-            r3, s3 = _print_expression(ns, node.operands[2])
-            if s2 and not s3:
-                r3 = "$signed({1'd0, " + r3 + "})"
-            if s3 and not s2:
-                r2 = "$signed({1'd0, " + r2 + "})"
-            r = r1 + " ? " + r2 + " : " + r3
-            s = s2 or s3
-        else:
-            raise TypeError
-        return "(" + r + ")", s
+        return _print_operator(ns, node)
 
     # Slice.
     elif isinstance(node, _Slice):
-        # Verilog does not like us slicing non-array signals...
-        if (isinstance(node.value, Signal) and
-            len(node.value) == 1 and
-            node.start == 0 and node.stop == 1):
-              return _print_expression(ns, node.value)
-
-        if node.start + 1 == node.stop:
-            sr = "[" + str(node.start) + "]"
-        else:
-            sr = "[" + str(node.stop-1) + ":" + str(node.start) + "]"
-        r, s = _print_expression(ns, node.value)
-        return r + sr, s
+        return _print_slice(ns, node)
 
     # Cat.
     elif isinstance(node, Cat):