log -header "Should be turned into a tree"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire a,
  input wire b,
  input wire c,
  input wire d,
  output wire x,
);
  assign x = a & b & c & d;
endmodule
EOF
check -assert

# Check equivalence after opt_balance_tree
equiv_opt -assert opt_balance_tree

design -load postopt

# Checks if inputs to and gates has been rewired
select -set driven_by_a i:a %co %co
select -set and_a_cell t:$and @driven_by_a %i

select -set driven_by_b i:b %co %co
select -set and_b_cell t:$and @driven_by_b %i

select -assert-none @and_a_cell @and_b_cell %d

select -set driven_by_c i:c %co %co
select -set and_c_cell t:$and @driven_by_c %i

select -set driven_by_d i:d %co %co
select -set and_d_cell t:$and @driven_by_d %i

select -assert-none @and_c_cell @and_d_cell %d

design -reset
log -pop

log -header "Should not be turned into a tree"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire a,
  input wire b,
  input wire c,
  input wire d,
  output wire x,
  output wire y,
  output wire z
);
  assign x = a & b;
  assign y = x & c;
  assign z = y & d;
endmodule
EOF
check -assert

# Check equivalence after opt_balance_tree
equiv_opt -assert opt_balance_tree

design -load postopt

# Checks if y is still wired up to the correct gate
select -set y_driver o:y %ci %ci
select -set and_y_cell t:$and @y_driver %i
select @and_y_cell -assert-count 1
select -set inputs @and_y_cell %ci
select -assert-count 1 @inputs i:c %i

design -reset
log -pop

log -header "With a cell"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire a,
  input wire b,
  input wire c,
  input wire d,
  output wire x,
  output wire y,
  output wire z
);
  wire temp;
  assign y = ~temp;
  assign x = a & b;
  assign temp = x & c;
  assign z = temp & d;
endmodule
EOF
check -assert

# Check equivalence after opt_balance_tree
equiv_opt -assert opt_balance_tree

design -load postopt
select -assert-count 3 t:$and

# Checks if temp is still wired up to the correct gate
select -set temp_driver w:temp %ci %ci
select -set and_cell t:$and @temp_driver %i
select @and_cell -assert-count 1
select -set inputs @and_cell %ci
select -assert-count 1 @inputs i:c %i

design -reset
log -pop

log -header "Word out port"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire a,
  input wire b,
  input wire c,
  input wire d,
  output wire [2:0] x
);
  assign x[0] = a & b;
  assign x[1] = x[0] & c;
  assign x[2] = x[1] & d;
endmodule
EOF
check -assert

# Check equivalence after opt_balance_tree
equiv_opt -assert opt_balance_tree

design -load postopt
select -assert-count 3 t:$and

# Checks if x[1] is still wired up to the correct gate
select -set target_drivers o:x %ci %ci
select -set target_cells t:$and @target_drivers %i
select -set inputs @target_cells %ci
select -assert-count 1 @inputs i:c %i

design -reset
log -pop

log -header "Fanout going to multiple outputs"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire a,
  input wire b,
  input wire c,
  input wire d,
  output wire [2:0] x,
  output wire y
);
  assign x[0] = a & b;
  assign x[1] = x[0] & c;
  assign x[2] = x[1] & d;

  assign y = x[1];
endmodule
EOF
check -assert

# Check equivalence after opt_balance_tree
equiv_opt -assert opt_balance_tree

design -load postopt
select -assert-count 3 t:$and

design -reset
log -pop

log -header "Fanout going to multiple outputs"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire a,
  input wire b,
  input wire c,
  input wire d,
  output wire [2:0] x,
  output wire y
);
  assign x[0] = a & b;
  assign x[1] = x[0] & c;
  assign x[2] = x[1] & d;

  assign y = x[1];
endmodule
EOF
check -assert

# Check equivalence after opt_balance_tree
equiv_opt -assert opt_balance_tree

design -load postopt
select -assert-count 3 t:$and

# Checks if y is still wired up to the correct gate
select -set y_wires o:y %ci
select -set y_driver @y_wires %ci %ci
select -set and_y_cell t:$and @y_driver %i
select -set inputs @and_y_cell %ci
select -assert-count 1 @inputs i:c %i

design -reset
log -pop

log -header "Interesting tree situation"
log -push
design -reset
read_verilog <<EOF
module top (
  input wire a, b, c, d,
  input wire e, f, g, h,
  output wire x,
);
  wire i, j;

  assign i = a & b & c & d;
  assign j = e & f & g & h & i;

  assign x = i & j;
endmodule
EOF
check -assert

# Check equivalence after opt_balance_tree
equiv_opt -assert opt_balance_tree

design -load postopt

# Checks if inputs to and gates has been rewired

# a and b
select -set driven_by_a i:a %co %co
select -set and_a_cell t:$and @driven_by_a %i

select -set driven_by_b i:b %co %co
select -set and_b_cell t:$and @driven_by_b %i

select -assert-none @and_a_cell @and_b_cell %d

# c and d
select -set driven_by_c i:c %co %co
select -set and_c_cell t:$and @driven_by_c %i

select -set driven_by_d i:d %co %co
select -set and_d_cell t:$and @driven_by_d %i

select -assert-none @and_c_cell @and_d_cell %d

# e and f
select -set driven_by_e i:e %co %co
select -set and_e_cell t:$and @driven_by_e %i

select -set driven_by_f i:f %co %co
select -set and_f_cell t:$and @driven_by_f %i

select -assert-none @and_e_cell @and_f_cell %d

# g and h
select -set driven_by_g i:g %co %co
select -set and_g_cell t:$and @driven_by_g %i

select -set driven_by_h i:h %co %co
select -set and_h_cell t:$and @driven_by_h %i

select -assert-none @and_g_cell @and_h_cell %d

design -reset
log -pop