From 5c277c6325b78bfe18cf294b63ea69ff272e69c5 Mon Sep 17 00:00:00 2001
From: Eddie Hung <eddie@fpgeh.com>
Date: Thu, 6 Jun 2019 14:21:34 -0700
Subject: [PATCH] Fix and test for balanced case

---
 passes/opt/muxpack.cc    | 24 ++++++++++++++----------
 tests/various/muxpack.v  | 26 ++++++++++++++++++++++++++
 tests/various/muxpack.ys | 15 +++++++++++++++
 3 files changed, 55 insertions(+), 10 deletions(-)

diff --git a/passes/opt/muxpack.cc b/passes/opt/muxpack.cc
index f9e5c8f09..8c4db4e4d 100644
--- a/passes/opt/muxpack.cc
+++ b/passes/opt/muxpack.cc
@@ -94,23 +94,27 @@ struct MuxpackWorker
 		{
 			log_debug("Considering %s (%s)\n", log_id(cell), log_id(cell->type));
 
-			SigSpec next_sig = cell->getPort("\\A");
-			if (sig_chain_prev.count(next_sig) == 0) {
-				if (cell->type == "$mux") {
-					next_sig = cell->getPort("\\B");
-					if (sig_chain_prev.count(next_sig) == 0)
-						goto start_cell;
-				}
-				else
+			SigSpec a_sig = cell->getPort("\\A");
+			if (cell->type == "$mux") {
+				SigSpec b_sig = cell->getPort("\\B");
+				if (sig_chain_prev.count(a_sig) + sig_chain_prev.count(b_sig) != 1)
+					goto start_cell;
+
+				if (!sig_chain_prev.count(a_sig))
+					a_sig = b_sig;
+			}
+			else if (cell->type == "$pmux") {
+				if (!sig_chain_prev.count(a_sig))
 					goto start_cell;
 			}
+			else log_abort();
 
 			{
-				for (auto bit : next_sig.bits())
+				for (auto bit : a_sig.bits())
 					if (sigbit_with_non_chain_users.count(bit))
 						goto start_cell;
 
-				Cell *c1 = sig_chain_prev.at(next_sig);
+				Cell *c1 = sig_chain_prev.at(a_sig);
 				Cell *c2 = cell;
 
 				if (c1->getParam("\\WIDTH") != c2->getParam("\\WIDTH"))
diff --git a/tests/various/muxpack.v b/tests/various/muxpack.v
index 7c189fff8..f1bd5ea8e 100644
--- a/tests/various/muxpack.v
+++ b/tests/various/muxpack.v
@@ -110,3 +110,29 @@ always @* begin
     endcase
 end
 endmodule
+
+module mux_if_bal_8_2 #(parameter N=8, parameter W=2) (input [N*W-1:0] i, input [$clog2(N)-1:0] s, output reg [W-1:0] o);
+always @*
+    if (s[0] == 1'b0)
+     if (s[1] == 1'b0)
+      if (s[2] == 1'b0)
+       o <= i[0*W+:W];
+      else
+       o <= i[1*W+:W];
+     else
+      if (s[2] == 1'b0)
+       o <= i[2*W+:W];
+      else
+       o <= i[3*W+:W];
+    else
+     if (s[1] == 1'b0)
+      if (s[2] == 1'b0)
+       o <= i[4*W+:W];
+      else
+       o <= i[5*W+:W];
+     else
+      if (s[2] == 1'b0)
+       o <= i[6*W+:W];
+      else
+       o <= i[7*W+:W];
+endmodule
diff --git a/tests/various/muxpack.ys b/tests/various/muxpack.ys
index 0c5b82818..9ea743b9f 100644
--- a/tests/various/muxpack.ys
+++ b/tests/various/muxpack.ys
@@ -133,3 +133,18 @@ design -import gold -as gold
 design -import gate -as gate
 miter -equiv -flatten -make_assert -make_outputs gold gate miter
 sat -verify -prove-asserts -show-ports miter
+
+design -load read
+hierarchy -top mux_if_bal_8_2
+prep
+design -save gold
+muxpack
+opt
+stat
+select -assert-count 7 t:$mux
+select -assert-count 0 t:$pmux
+design -stash gate
+design -import gold -as gold
+design -import gate -as gate
+miter -equiv -flatten -make_assert -make_outputs gold gate miter
+sat -verify -prove-asserts -show-ports miter
-- 
2.30.2