log -header "Simple MUX chain to PMUX"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [2:0] sel,
  input wire [3:0] a,
  output wire x
);
  wire w0, w1;

  assign w0 = sel[0] ? a[1] : a[0];
  assign w1 = sel[1] ? a[2] : w0;
  assign x  = sel[2] ? a[3] : w1;
endmodule
EOF
check -assert

# Check equivalence after extract_reduce
equiv_opt -assert extract_reduce

# Load design and run opt_clean to remove unnecessary wires
design -load postopt
opt_clean

# Check we got a single pmux with the correct input number
select -assert-count 0 t:$mux
select -assert-count 1 t:$pmux
select -assert-count 1 t:$pmux r:S_WIDTH=4 %i

design -reset
log -pop



log -header "MUX chain with constants"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [2:0] sel,
  input wire [2:0] a,
  output wire x
);
  wire w0, w1;

  assign w0 = sel[0] ? a[1] : a[0];
  assign w1 = sel[1] ? a[2] : w0;
  assign x  = sel[2] ? 1'b1 : w1;
endmodule
EOF
check -assert

# Check equivalence after extract_reduce
equiv_opt -assert extract_reduce

# Load design and run opt_clean to remove unnecessary wires
design -load postopt
opt_clean

# Check we got a single pmux with the correct input number
select -assert-count 0 t:$mux
select -assert-count 1 t:$pmux
select -assert-count 1 t:$pmux r:S_WIDTH=4 %i

design -reset
log -pop



log -header "MUX chain with multiple branches"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [2:0] sel,
  input wire [3:0] a,
  output wire x
);
  wire w0, w1;

  assign w0 = sel[0] ? a[1] : a[0];
  assign w1 = sel[1] ? a[2] : a[3];
  assign x  = sel[2] ? w0 : w1;
endmodule
EOF
check -assert

# Check equivalence after extract_reduce
equiv_opt -assert extract_reduce

# Load design and run opt_clean to remove unnecessary wires
design -load postopt
opt_clean

# Check we got a single pmux with the correct input number
select -assert-count 0 t:$mux
select -assert-count 1 t:$pmux
select -assert-count 1 t:$pmux r:S_WIDTH=4 %i

design -reset
log -pop



log -header "MUX chain with multiple uneven branches"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [5:0] sel,
  input wire [6:0] a,
  output wire x
);
  wire w0, w1, w2, w3, w4;

  assign w0 = sel[0] ? a[1] : a[0];
  assign w1 = sel[1] ? a[2] : w0;
  assign w2 = sel[2] ? a[3] : w1;
  assign w3 = sel[3] ? w2 : w4;
  assign w4 = sel[5] ? a[4] : a[5];
  assign x  = sel[4] ? w3 : a[6];
endmodule
EOF
check -assert

# Check equivalence after extract_reduce
equiv_opt -assert extract_reduce

# Load design and run opt_clean to remove unnecessary wires
design -load postopt
opt_clean

# Check we got a single pmux with the correct input number
select -assert-count 0 t:$mux
select -assert-count 1 t:$pmux
select -assert-count 1 t:$pmux r:S_WIDTH=7 %i

design -reset
log -pop



log -header "No off-chain MUX chain"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [3:0] sel,
  input wire [4:0] a,
  output wire x,
  output wire y
);
  wire w0, w1, w2;

  assign w0 = sel[0] ? a[1] : a[0];
  assign w1 = sel[1] ? a[2] : w0;
  assign w2 = sel[2] ? a[3] : w1;
  assign x  = sel[3] ? a[4] : w2;

  // Off-chain use of intermediate wire
  assign y = w1;
endmodule
EOF
check -assert

# Check equivalence after extract_reduce
equiv_opt -assert extract_reduce

# Load design and run opt_clean to remove unnecessary wires
design -load postopt
opt_clean

# Check we got two pmuxes
select -assert-count 0 t:$mux
select -assert-count 2 t:$pmux

# Check that both pmuxes have input width of 3
select -assert-none t:$pmux r:S_WIDTH!=3 %i

design -reset
log -pop




log -header "Allow off-chain MUX chain"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [3:0] sel,
  input wire [4:0] a,
  output wire x,
  output wire y
);
  wire w0, w1, w2;

  assign w0 = sel[0] ? a[1] : a[0];
  assign w1 = sel[1] ? a[2] : w0;
  assign w2 = sel[2] ? a[3] : w1;
  assign x  = sel[3] ? a[4] : w2;

  assign y = w1;
endmodule
EOF
check -assert

# Check equivalence after extract_reduce
equiv_opt -assert extract_reduce -allow-off-chain

# Load design and run opt_clean to remove unnecessary wires
design -load postopt
opt_clean

# Check we got two pmux
select -assert-count 0 t:$mux
select -assert-count 2 t:$pmux

# Check that one pmux has an input width of 3 
# and the other has an input width of 5
select -assert-count 1 t:$pmux r:S_WIDTH=3 %i
select -assert-count 1 t:$pmux r:S_WIDTH=5 %i

design -reset
log -pop



log -header "Reconverging tree; no off-chain"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [7:0] sel,
  input wire [7:0] a,
  output wire x
);
  wire w0, w1, w2, w3, w4, w5, w6;

  assign w0 = sel[0] ? a[1] : a[0];
  assign w1 = sel[1] ? a[2] : w0;
  assign w2 = sel[2] ? w4 : w1;
  assign w3 = sel[3] ? w2 : w5;
  assign w4 = sel[6] ? a[3] : w6;
  assign w5 = sel[5] ? a[5] : w4;
  assign w6 = sel[7] ? a[7] : a[4];
  assign x  = sel[4] ? a[6] : w3;
endmodule
EOF
check -assert

# Check equivalence after extract_reduce
equiv_opt -assert extract_reduce

# Load design and run opt_clean and opt_reduce to remove unnecessary cells
design -load postopt
opt_clean
opt_reduce

# Check we got one pmux with correct number of inputs
select -assert-count 0 t:$mux
select -assert-count 2 t:$pmux
select -assert-count 1 t:$pmux r:S_WIDTH=3 %i
select -assert-count 1 t:$pmux r:S_WIDTH=6 %i

design -reset
log -pop



# log -header "Reconverging tree; yes off-chain"
# log -push
# design -reset
# read_verilog <<EOF
# module top (
#   input wire [7:0] sel,
#   input wire [7:0] a,
#   output wire x
# );
#   wire w0, w1, w2, w3, w4, w5, w6;

#   assign w0 = sel[0] ? a[1] : a[0];
#   assign w1 = sel[1] ? a[2] : w0;
#   assign w2 = sel[2] ? w4 : w1;
#   assign w3 = sel[3] ? w2 : w5;
#   assign w4 = sel[6] ? a[3] : w6;
#   assign w5 = sel[5] ? a[5] : w4;
#   assign w6 = sel[7] ? a[7] : a[4];
#   assign x  = sel[4] ? a[6] : w3;
# endmodule
# EOF
# check -assert

# # Check equivalence after extract_reduce
# equiv_opt -assert extract_reduce -allow-off-chain

# # Load design and run opt_clean and opt_reduce to remove unnecessary cells
# design -load postopt
# opt_clean
# opt_reduce

# # Check we got one pmux with correct number of inputs
# select -assert-count 0 t:$mux
# select -assert-count 1 t:$pmux
# select -assert-count 1 t:$pmux r:S_WIDTH=8 %i

# design -reset
# log -pop



log -header "Reusing select inputs"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [5:0] sel,
  input wire [7:0] a,
  output wire x
);
  wire w0, w1, w2, w3, w4, w5, w6;

  assign w0 = sel[0] ? a[1] : a[0];
  assign w1 = sel[1] ? a[2] : w0;
  assign w2 = sel[2] ? w4 : w1;
  assign w3 = sel[3] ? w2 : w5;
  assign w4 = sel[2] ? a[3] : w6;
  assign w5 = sel[5] ? a[5] : w4;
  assign w6 = sel[4] & sel[1] ? a[7] : a[4];
  assign x  = sel[4] ? a[6] : w3;
endmodule
EOF
check -assert

# Check equivalence after extract_reduce
equiv_opt -assert extract_reduce

# Load design and run opt_clean and opt_reduce to remove unnecessary cells
design -load postopt
opt_clean
opt_reduce

# Check we got one pmux with correct number of inputs
select -assert-count 0 t:$mux
select -assert-count 2 t:$pmux
select -assert-count 1 t:$pmux r:S_WIDTH=3 %i
select -assert-count 1 t:$pmux r:S_WIDTH=6 %i

design -reset
log -pop



log -header "Using outputs as select inputs"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [5:0] sel,
  input wire [7:0] a,
  output wire x
);
  wire w0, w1, w2, w3, w4, w5, w6;

  assign w0 = sel[0] ? a[1] : a[0];
  assign w1 = sel[1] ? a[2] : w0;
  assign w2 = sel[2] ? w4 : w1;
  assign w3 = sel[3] ? w2 : w5;
  assign w4 = sel[2] ? a[3] : w6;
  assign w5 = sel[5] ? a[5] : w4;
  assign w6 = a[2] ? a[7] : a[4];
  assign x  = sel[4] ? a[6] : w3;
endmodule
EOF
check -assert

# Check equivalence after extract_reduce
equiv_opt -assert extract_reduce

# Load design and run opt_clean and opt_reduce to remove unnecessary cells
design -load postopt
opt_clean
opt_reduce

# Check we got one pmux with correct number of inputs
select -assert-count 0 t:$mux
select -assert-count 2 t:$pmux
select -assert-count 1 t:$pmux r:S_WIDTH=3 %i
select -assert-count 1 t:$pmux r:S_WIDTH=6 %i

design -reset
log -pop



log -header "Normal stress test"
log -push
design -reset
read_verilog <<EOF
module top (
  input  wire [5:0] sel0,
  input  wire [5:0] sel1,
  input  wire [5:0] sel2,
  input  wire [6:0] a0,
  input  wire [6:0] a1,
  input  wire [6:0] a2,
  output wire x
);
  wire x0, x1;

  // Stage 0
  wire s0_w0 = sel0[0] ? a0[1] : a0[0];
  wire s0_w1 = sel0[1] ? a0[2] : s0_w0;
  wire s0_w2 = sel0[2] ? a0[3] : s0_w1;
  wire s0_w4 = sel0[5] ? a0[4] : a0[5];
  wire s0_w3 = sel0[3] ? s0_w2 : s0_w4;
  assign x0  = sel0[4] ? s0_w3 : a0[6];

  // Stage 1
  wire s1_w0 = sel1[0] ? a1[1] : a1[0];
  wire s1_w1 = sel1[1] ? a1[2] : s1_w0;
  wire s1_w2 = sel1[2] ? a1[3] : s1_w1;
  wire s1_w4 = sel1[5] ? a1[4] : a1[5];
  wire s1_w3 = sel1[3] ? s1_w2 : s1_w4;
  assign x1  = sel1[4] ? s1_w3 : x0;

  // Stage 2
  wire s2_w0 = sel2[0] ? a2[1] : a2[0];
  wire s2_w1 = sel2[1] ? a2[2] : s2_w0;
  wire s2_w2 = sel2[2] ? a2[3] : s2_w1;
  wire s2_w4 = sel2[5] ? a2[4] : a2[5];
  wire s2_w3 = sel2[3] ? s2_w2 : s2_w4;
  assign x   = sel2[4] ? s2_w3 : x1;
endmodule
EOF
check -assert

# Check equivalence after extract_reduce
equiv_opt -assert extract_reduce

# Load design and run opt_clean and opt_reduce to remove unnecessary cells
design -load postopt
opt_clean
opt_reduce

# Check we got one pmux with correct number of inputs
select -assert-count 0 t:$mux
select -assert-count 1 t:$pmux

design -reset
log -pop



log -header "Stress test with reconverging tree and no off chain"
log -push
design -reset
read_verilog <<EOF
module top (
  input  wire [7:0] sel0,
  input  wire [7:0] sel1,
  input  wire [7:0] sel2,
  input  wire [7:0] a0,
  input  wire [7:0] a1,
  input  wire [7:0] a2,
  output wire x
);
  wire x0, x1;

  // Stage 0
  wire s0_w6 = sel0[7] ? a0[7] : a0[4];
  wire s0_w4 = sel0[6] ? a0[3] : s0_w6;
  wire s0_w5 = sel0[5] ? a0[5] : s0_w4;
  wire s0_w0 = sel0[0] ? a0[1] : a0[0];
  wire s0_w1 = sel0[1] ? a0[2] : s0_w0;
  wire s0_w2 = sel0[2] ? s0_w4 : s0_w1;
  wire s0_w3 = sel0[3] ? s0_w2 : s0_w5;
  assign x0  = sel0[4] ? a0[6] : s0_w3;

  // Stage 1
  wire s1_w6 = sel1[7] ? a1[7] : a1[4];
  wire s1_w4 = sel1[6] ? a1[3] : s1_w6;
  wire s1_w5 = sel1[5] ? a1[5] : s1_w4;
  wire s1_w0 = sel1[0] ? a1[1] : x0;
  wire s1_w1 = sel1[1] ? a1[2] : s1_w0;
  wire s1_w2 = sel1[2] ? s1_w4 : s1_w1;
  wire s1_w3 = sel1[3] ? s1_w2 : s1_w5;
  assign x1  = sel1[4] ? a1[6] : s1_w3;

  // Stage 2
  wire s2_w6 = sel2[7] ? a2[7] : a2[4];
  wire s2_w4 = sel2[6] ? a2[3] : s2_w6;
  wire s2_w5 = sel2[5] ? a2[5] : s2_w4;
  wire s2_w0 = sel2[0] ? a2[1] : x1;
  wire s2_w1 = sel2[1] ? a2[2] : s2_w0;
  wire s2_w2 = sel2[2] ? s2_w4 : s2_w1;
  wire s2_w3 = sel2[3] ? s2_w2 : s2_w5;
  assign x   = sel2[4] ? a2[6] : s2_w3;
endmodule

EOF
check -assert

# Check equivalence after extract_reduce
equiv_opt -assert extract_reduce

# Load design and run opt_clean and opt_reduce to remove unnecessary cells
design -load postopt
opt_clean
opt_reduce

# Check we got one pmux with correct number of inputs
select -assert-count 0 t:$mux
select -assert-count 4 t:$pmux

design -reset
log -pop



# log -header "Stress test with reconverging tree and yes off chain"
# log -push
# design -reset
# read_verilog <<EOF
# module top (
#   input  wire [7:0] sel0,
#   input  wire [7:0] sel1,
#   input  wire [7:0] sel2,
#   input  wire [7:0] a0,
#   input  wire [7:0] a1,
#   input  wire [7:0] a2,
#   output wire x
# );
#   wire x0, x1;

#   // Stage 0
#   wire s0_w6 = sel0[7] ? a0[7] : a0[4];
#   wire s0_w4 = sel0[6] ? a0[3] : s0_w6;
#   wire s0_w5 = sel0[5] ? a0[5] : s0_w4;
#   wire s0_w0 = sel0[0] ? a0[1] : a0[0];
#   wire s0_w1 = sel0[1] ? a0[2] : s0_w0;
#   wire s0_w2 = sel0[2] ? s0_w4 : s0_w1;
#   wire s0_w3 = sel0[3] ? s0_w2 : s0_w5;
#   assign x0  = sel0[4] ? a0[6] : s0_w3;

#   // Stage 1
#   wire s1_w6 = sel1[7] ? a1[7] : a1[4];
#   wire s1_w4 = sel1[6] ? a1[3] : s1_w6;
#   wire s1_w5 = sel1[5] ? a1[5] : s1_w4;
#   wire s1_w0 = sel1[0] ? a1[1] : x0;
#   wire s1_w1 = sel1[1] ? a1[2] : s1_w0;
#   wire s1_w2 = sel1[2] ? s1_w4 : s1_w1;
#   wire s1_w3 = sel1[3] ? s1_w2 : s1_w5;
#   assign x1  = sel1[4] ? a1[6] : s1_w3;

#   // Stage 2
#   wire s2_w6 = sel2[7] ? a2[7] : a2[4];
#   wire s2_w4 = sel2[6] ? a2[3] : s2_w6;
#   wire s2_w5 = sel2[5] ? a2[5] : s2_w4;
#   wire s2_w0 = sel2[0] ? a2[1] : x1;
#   wire s2_w1 = sel2[1] ? a2[2] : s2_w0;
#   wire s2_w2 = sel2[2] ? s2_w4 : s2_w1;
#   wire s2_w3 = sel2[3] ? s2_w2 : s2_w5;
#   assign x   = sel2[4] ? a2[6] : s2_w3;
# endmodule

# EOF
# check -assert

# # Check equivalence after extract_reduce
# equiv_opt -assert extract_reduce -allow-off-chain

# # Load design and run opt_clean and opt_reduce to remove unnecessary cells
# design -load postopt
# opt_clean
# opt_reduce

# # Check we got one pmux with correct number of inputs
# select -assert-count 0 t:$mux
# select -assert-count 1 t:$pmux

# design -reset
# log -pop