###################################################################
# Extract Reduce AND Gates Tests
###################################################################

log -header "Simple AND chain"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [3:0] a,
  output wire x
);
  assign x = a[0] & a[1] & a[2] & a[3];
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 final design has correct number of gates and inputs
select -assert-count 0 t:$and
select -assert-count 1 t:$reduce_and r:A_WIDTH=4 %i

design -reset
log -pop



log -header "AND chain with constants"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [2:0] a,
  output wire x
);
  assign x = a[0] & a[1] & a[2] & 1'b1;
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 final design has correct number of gates and inputs
select -assert-count 0 t:$and
select -assert-count 1 t:$reduce_and r:A_WIDTH=4 %i

design -reset
log -pop



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

  assign w0 = a[0] & a[1];
  assign w1 = a[2] & a[3];
  assign w2 = a[4] & a[5];
  assign w3 = w0 & w1;
  assign x = w2 & w3;
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 final design has correct number of gates and inputs
select -assert-count 0 t:$and
select -assert-count 1 t:$reduce_and r:A_WIDTH=6 %i

design -reset
log -pop



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

  assign w0 = a[0] & a[1];
  assign w1 = w0 & a[2];
  assign w2 = w1 & a[3];
  assign x = w2 & a[4];

  // Off-chain use of w1
  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 final design has correct number of gates
select -assert-count 0 t:$and
select -assert-count 2 t:$reduce_and

# Check that both gates are 3 bits wide
select -assert-none t:$reduce_and r:A_WIDTH!=3 %i

design -reset
log -pop



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

  assign w0 = a[0] & a[1];
  assign w1 = w0 & a[2];
  assign w2 = w1 & a[3];
  assign x = w2 & a[4];

  // Off-chain use of w1
  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 final design has correct number of gates
select -assert-count 0 t:$and
select -assert-count 2 t:$reduce_and

# Check that only one gate has a width of 5 and one gate has a width of 3
select -assert-count 1 t:$reduce_and r:A_WIDTH=5 %i
select -assert-count 1 t:$reduce_and r:A_WIDTH=3 %i

design -reset
log -pop



log -header "No off-chain with branches for AND"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [10:0] a,
  output wire x,
  output wire y
);
  wire w0;

  assign w0 = a[0] & a[1] & a[2] & (a[3] & a[4]) & a[5];
  assign x  = w0 & a[6] & a[7] & (a[8] & a[9]) & a[10];

  // Off-chain use of w0
  assign y = w0;
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 final design has correct number of gates
select -assert-count 0 t:$and
select -assert-count 2 t:$reduce_and

# Check that both gates are 6 bits wide
select -assert-none t:$reduce_and r:A_WIDTH!=6 %i

design -reset
log -pop



log -header "Allow off-chain with branches for AND"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [10:0] a,
  output wire x,
  output wire y
);
  wire w0;

  assign w0 = a[0] & a[1] & a[2] & (a[3] & a[4]) & a[5];
  assign x  = w0 & a[6] & a[7] & (a[8] & a[9]) & a[10];

  // Off-chain use of w0
  assign y = w0;
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 final design has correct number of gates
select -assert-count 0 t:$and
select -assert-count 2 t:$reduce_and

# Check that only one gate has a width of 11 and one gate has a width of 6
select -assert-count 1 t:$reduce_and r:A_WIDTH=11 %i
select -assert-count 1 t:$reduce_and r:A_WIDTH=6 %i

design -reset
log -pop

###################################################################
# Extract Reduce OR Gates Tests
###################################################################

log -header "Simple OR chain"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [3:0] a,
  output wire x
);
  assign x = a[0] | a[1] | a[2] | a[3];
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 final design has correct number of gates and inputs
select -assert-count 0 t:$or
select -assert-count 1 t:$reduce_or r:A_WIDTH=4 %i

design -reset
log -pop



log -header "OR chain with constants"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [3:0] a,
  output wire x
);
  assign x = a[0] | a[1] | a[2] | 1'b1;
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 final design has correct number of gates and inputs
select -assert-count 0 t:$or
select -assert-count 1 t:$reduce_or r:A_WIDTH=4 %i

design -reset
log -pop



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

  assign w0 = a[0] | a[1];
  assign w1 = a[2] | a[3];
  assign w2 = a[4] | a[5];
  assign w3 = w0 | w1;
  assign x = w2 | w3;
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 final design has correct number of gates and inputs
select -assert-count 0 t:$or
select -assert-count 1 t:$reduce_or r:A_WIDTH=6 %i

design -reset
log -pop



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

  assign w0 = a[0] | a[1];
  assign w1 = w0 | a[2];
  assign w2 = w1 | a[3];
  assign x = w2 | a[4];

  // Off-chain use of w1
  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 final design has correct number of gates
select -assert-count 0 t:$or
select -assert-count 2 t:$reduce_or

# Check that both gates are 3 bits wide
select -assert-none t:$reduce_or r:A_WIDTH!=3 %i

design -reset
log -pop



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

  assign w0 = a[0] | a[1];
  assign w1 = w0 | a[2];
  assign w2 = w1 | a[3];
  assign x = w2 | a[4];

  // Off-chain use of w1
  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 final design has correct number of gates
select -assert-count 0 t:$or
select -assert-count 2 t:$reduce_or

# Check that only one gate has a width of 5 and one gate has a width of 3
select -assert-count 1 t:$reduce_or r:A_WIDTH=5 %i
select -assert-count 1 t:$reduce_or r:A_WIDTH=3 %i

design -reset
log -pop



log -header "No off-chain with branches for OR"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [10:0] a,
  output wire x,
  output wire y
);
  wire w0;

  assign w0 = a[0] | a[1] | a[2] | (a[3] | a[4]) | a[5];
  assign x  = w0 | a[6] | a[7] | (a[8] | a[9]) | a[10];

  // Off-chain use of w0
  assign y = w0;
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 final design has correct number of gates
select -assert-count 0 t:$and
select -assert-count 2 t:$reduce_or

# Check that both gates are 6 bits wide
select -assert-none t:$reduce_or r:A_WIDTH!=6 %i

design -reset
log -pop



log -header "Allow off-chain with branches for OR"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [10:0] a,
  output wire x,
  output wire y
);
  wire w0;

  assign w0 = a[0] | a[1] | a[2] | (a[3] | a[4]) | a[5];
  assign x  = w0 | a[6] | a[7] | (a[8] | a[9]) | a[10];

  // Off-chain use of w0
  assign y = w0;
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 final design has correct number of gates
select -assert-count 0 t:$and
select -assert-count 2 t:$reduce_or

# Check that only one gate has a width of 11 and one gate has a width of 6
select -assert-count 1 t:$reduce_or r:A_WIDTH=11 %i
select -assert-count 1 t:$reduce_or r:A_WIDTH=6 %i

design -reset
log -pop

###################################################################
# Extract Reduce XOR Gates Tests
###################################################################

log -header "Simple XOR chain"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [3:0] a,
  output wire x
);
  assign x = a[0] ^ a[1] ^ a[2] ^ a[3];
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 final design has correct number of gates and inputs
select -assert-count 0 t:$xor
select -assert-count 1 t:$reduce_xor r:A_WIDTH=4 %i

design -reset
log -pop



log -header "XOR chain with constants"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [3:0] a,
  output wire x
);
  assign x = a[0] ^ a[1] ^ a[2] ^ 1'b1;
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 final design has correct number of gates and inputs
select -assert-count 0 t:$xor
select -assert-count 1 t:$reduce_xor r:A_WIDTH=4 %i

design -reset
log -pop



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

  assign w0 = a[0] ^ a[1];
  assign w1 = a[2] ^ a[3];
  assign w2 = a[4] ^ a[5];
  assign w3 = w0 ^ w1;
  assign x = w2 ^ w3;
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 final design has correct number of gates and inputs
select -assert-count 0 t:$xor
select -assert-count 1 t:$reduce_xor r:A_WIDTH=6 %i

design -reset
log -pop



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

  assign w0 = a[0] ^ a[1];
  assign w1 = w0 ^ a[2];
  assign w2 = w1 ^ a[3];
  assign x = w2 ^ a[4];

  // Off-chain use of w1
  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 final design has correct number of gates
select -assert-count 0 t:$xor
select -assert-count 2 t:$reduce_xor

# Check that both gates are 3 bits wide
select -assert-none t:$reduce_xor r:A_WIDTH!=3 %i

design -reset
log -pop



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

  assign w0 = a[0] ^ a[1];
  assign w1 = w0 ^ a[2];
  assign w2 = w1 ^ a[3];
  assign x = w2 ^ a[4];

  // Off-chain use of w1
  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 final design has correct number of gates
select -assert-count 0 t:$xor
select -assert-count 2 t:$reduce_xor

# Check that only one gate has a width of 5 and one gate has a width of 3
select -assert-count 1 t:$reduce_xor r:A_WIDTH=5 %i
select -assert-count 1 t:$reduce_xor r:A_WIDTH=3 %i

design -reset
log -pop



log -header "No off-chain with branches for XOR"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [10:0] a,
  output wire x,
  output wire y
);
  wire w0;

  assign w0 = a[0] ^ a[1] ^ a[2] ^ (a[3] ^ a[4]) ^ a[5];
  assign x  = w0 ^ a[6] ^ a[7] ^ (a[8] ^ a[9]) ^ a[10];

  // Off-chain use of w0
  assign y = w0;
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 final design has correct number of gates
select -assert-count 0 t:$xor
select -assert-count 2 t:$reduce_xor

# Check that both gates are 6 bits wide
select -assert-none t:$reduce_xor r:A_WIDTH!=6 %i

design -reset
log -pop



log -header "Allow off-chain with branches for XOR"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire [10:0] a,
  output wire x,
  output wire y
);
  wire w0;

  assign w0 = a[0] ^ a[1] ^ a[2] ^ (a[3] ^ a[4]) ^ a[5];
  assign x  = w0 ^ a[6] ^ a[7] ^ (a[8] ^ a[9]) ^ a[10];

  // Off-chain use of w0
  assign y = w0;
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 final design has correct number of gates
select -assert-count 0 t:$xor
select -assert-count 2 t:$reduce_xor

# Check that only one gate has a width of 11 and one gate has a width of 6
select -assert-count 1 t:$reduce_xor r:A_WIDTH=11 %i
select -assert-count 1 t:$reduce_xor r:A_WIDTH=6 %i

design -reset
log -pop

###################################################################
# Extract PMUX from MUX Chains Tests
###################################################################

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 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 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 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 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