From d77beecc6e957d2b82189a4ed3d103cb724d1855 Mon Sep 17 00:00:00 2001
From: Wilson Snyder <wsnyder@wsnyder.org>
Date: Fri, 27 Nov 2020 21:34:40 -0500
Subject: [PATCH] Fix $countbits(..., 'z)

---
 src/V3Number.cpp                  |  5 +--
 src/V3Number.h                    |  1 +
 src/V3Tristate.cpp                | 59 +++++++++++++++++++++++++++++++
 test_regress/t/t_math_countbits.v |  2 --
 4 files changed, 63 insertions(+), 4 deletions(-)

diff --git a/src/V3Number.cpp b/src/V3Number.cpp
index 8d4e57307..a2073dff7 100644
--- a/src/V3Number.cpp
+++ b/src/V3Number.cpp
@@ -937,10 +937,11 @@ bool V3Number::isAnyX() const {
     }
     return false;
 }
-bool V3Number::isAnyXZ() const {
+bool V3Number::isAnyXZ() const { return isAnyX() || isAnyZ(); }
+bool V3Number::isAnyZ() const {
     if (isDouble() || isString()) return false;
     for (int bit = 0; bit < width(); bit++) {
-        if (bitIsX(bit) || bitIsZ(bit)) return true;
+        if (bitIsZ(bit)) return true;
     }
     return false;
 }
diff --git a/src/V3Number.h b/src/V3Number.h
index 7a01876fb..240858ea9 100644
--- a/src/V3Number.h
+++ b/src/V3Number.h
@@ -273,6 +273,7 @@ public:
     void isSigned(bool ssigned) { m_signed = ssigned; }
     bool isAnyX() const;
     bool isAnyXZ() const;
+    bool isAnyZ() const;
     bool isMsbXZ() const { return bitIsXZ(m_width); }
     uint32_t toUInt() const;
     vlsint32_t toSInt() const;
diff --git a/src/V3Tristate.cpp b/src/V3Tristate.cpp
index ca5119846..cc192edd4 100644
--- a/src/V3Tristate.cpp
+++ b/src/V3Tristate.cpp
@@ -981,6 +981,65 @@ class TristateVisitor final : public TristateBaseVisitor {
     virtual void visit(AstEqWild* nodep) override { visitEqNeqWild(nodep); }
     virtual void visit(AstNeqWild* nodep) override { visitEqNeqWild(nodep); }
 
+    virtual void visit(AstCountBits* nodep) override {
+        std::array<bool, 3> dropop;
+        dropop[0] = VN_IS(nodep->rhsp(), Const) && VN_CAST(nodep->rhsp(), Const)->num().isAnyZ();
+        dropop[1] = VN_IS(nodep->thsp(), Const) && VN_CAST(nodep->thsp(), Const)->num().isAnyZ();
+        dropop[2] = VN_IS(nodep->fhsp(), Const) && VN_CAST(nodep->fhsp(), Const)->num().isAnyZ();
+        UINFO(4, " COUNTBITS(" << dropop[0] << dropop[1] << dropop[2] << " " << nodep << endl);
+        AstVarRef* varrefp = VN_CAST(nodep->lhsp(), VarRef);  // Input variable
+        if (m_graphing) {
+            iterateAndNextNull(nodep->lhsp());
+            if (!dropop[0]) iterateAndNextNull(nodep->rhsp());
+            if (!dropop[1]) iterateAndNextNull(nodep->thsp());
+            if (!dropop[2]) iterateAndNextNull(nodep->fhsp());
+        } else {
+            AstNode* nonXp = nullptr;
+            if (!dropop[0])
+                nonXp = nodep->rhsp();
+            else if (!dropop[1])
+                nonXp = nodep->thsp();
+            else if (!dropop[2])
+                nonXp = nodep->fhsp();
+            // Replace 'z with non-Z
+            if (dropop[0] || dropop[1] || dropop[2]) {
+                // Unsupported: A $countones('0) should compare with the enables, but we don't
+                // do so at present, we only compare if there is a z in the equation. Otherwise
+                // we'd need to attach an enable to every signal, then optimize them away later
+                // when we determine the signal has no tristate
+                if (!VN_IS(nodep->lhsp(), VarRef)) {
+                    nodep->v3warn(E_UNSUPPORTED, "Unsupported LHS tristate construct: "
+                                                     << nodep->prettyTypeName());
+                    return;
+                }
+                AstVar* envarp = getCreateEnVarp(varrefp->varp());
+                // If any drops, we need to add in the count of Zs (from __en)
+                UINFO(4, " COUNTBITS('z)-> " << nodep << endl);
+                AstNRelinker relinkHandle;
+                nodep->unlinkFrBack(&relinkHandle);
+                AstNode* newp = new AstCountOnes(
+                    nodep->fileline(), new AstVarRef(nodep->fileline(), envarp, VAccess::READ));
+                if (nonXp) {  // Need to still count '0 or '1 or 'x's
+                    if (dropop[0]) {
+                        nodep->rhsp()->unlinkFrBack()->deleteTree();
+                        nodep->rhsp(nonXp->cloneTree(true));
+                    }
+                    if (dropop[1]) {
+                        nodep->thsp()->unlinkFrBack()->deleteTree();
+                        nodep->thsp(nonXp->cloneTree(true));
+                    }
+                    if (dropop[2]) {
+                        nodep->fhsp()->unlinkFrBack()->deleteTree();
+                        nodep->fhsp(nonXp->cloneTree(true));
+                    }
+                    newp = new AstAdd(nodep->fileline(), nodep, newp);
+                }
+                if (debug() >= 9) newp->dumpTree(cout, "-countout: ");
+                relinkHandle.relink(newp);
+            }
+            iterateChildren(nodep);
+        }
+    }
     virtual void visit(AstPull* nodep) override {
         UINFO(9, dbgState() << nodep << endl);
         AstVarRef* varrefp = nullptr;
diff --git a/test_regress/t/t_math_countbits.v b/test_regress/t/t_math_countbits.v
index de02ccb5c..547263faa 100644
--- a/test_regress/t/t_math_countbits.v
+++ b/test_regress/t/t_math_countbits.v
@@ -96,7 +96,6 @@ module t(/*AUTOARG*/
          if ($countbits(val, '1, '1, 'x) != 7) $stop;
          if ($countbits(val, 'x, '0) != 25) $stop;
          if ($countbits(val, 'x, '0, '1) != 32) $stop;
-`ifndef VERILATOR
          // Optimization may depend on position of Z, so need to walk it
          if ($countbits(val, 'z) != 0) $stop;
          if ($countbits(val, 'z, '1) != 7) $stop;
@@ -106,7 +105,6 @@ module t(/*AUTOARG*/
          if ($countbits(val, 'z, '0, '1) != 32) $stop;
          //
          if ($countbits(val, 'x, 'z) != 0) $stop;
-`endif
       end
       else if (cyc == 1) begin
          in16 <= 16'h0AF0;