@@ -50,31 +50,36 @@ impl MergeExpressionsOptimizer {
50
50
let mut new_circuit = Vec :: new ( ) ;
51
51
let mut new_acir_opcode_positions = Vec :: new ( ) ;
52
52
// For each opcode, try to get a target opcode to merge with
53
- for ( i, opcode) in circuit. opcodes . iter ( ) . enumerate ( ) {
53
+ for ( i, ( opcode, opcode_position) ) in
54
+ circuit. opcodes . iter ( ) . zip ( acir_opcode_positions) . enumerate ( )
55
+ {
54
56
if !matches ! ( opcode, Opcode :: AssertZero ( _) ) {
55
57
new_circuit. push ( opcode. clone ( ) ) ;
56
- new_acir_opcode_positions. push ( acir_opcode_positions [ i ] ) ;
58
+ new_acir_opcode_positions. push ( opcode_position ) ;
57
59
continue ;
58
60
}
59
61
let opcode = modified_gates. get ( & i) . unwrap_or ( opcode) . clone ( ) ;
60
62
let mut to_keep = true ;
61
63
let input_witnesses = self . witness_inputs ( & opcode) ;
62
- for w in input_witnesses. clone ( ) {
63
- let empty_gates = BTreeSet :: new ( ) ;
64
- let gates_using_w = used_witness. get ( & w) . unwrap_or ( & empty_gates) ;
64
+ for w in input_witnesses {
65
+ let Some ( gates_using_w) = used_witness. get ( & w) else {
66
+ continue ;
67
+ } ;
65
68
// We only consider witness which are used in exactly two arithmetic gates
66
69
if gates_using_w. len ( ) == 2 {
67
- let gates_using_w : Vec < _ > = gates_using_w. iter ( ) . collect ( ) ;
68
- let mut b = * gates_using_w[ 1 ] ;
69
- if b == i {
70
- b = * gates_using_w [ 0 ] ;
70
+ let first = * gates_using_w. first ( ) . expect ( "gates_using_w.len == 2" ) ;
71
+ let second = * gates_using_w. last ( ) . expect ( "gates_using_w.len == 2" ) ;
72
+ let b = if second == i {
73
+ first
71
74
} else {
72
75
// sanity check
73
- assert ! ( i == * gates_using_w[ 0 ] ) ;
74
- }
75
- let second_gate = modified_gates. get ( & b) . unwrap_or ( & circuit. opcodes [ b] ) . clone ( ) ;
76
+ assert ! ( i == first) ;
77
+ second
78
+ } ;
79
+
80
+ let second_gate = modified_gates. get ( & b) . unwrap_or ( & circuit. opcodes [ b] ) ;
76
81
if let ( Opcode :: AssertZero ( expr_define) , Opcode :: AssertZero ( expr_use) ) =
77
- ( opcode. clone ( ) , second_gate)
82
+ ( & opcode, second_gate)
78
83
{
79
84
// We cannot merge an expression into an earlier opcode, because this
80
85
// would break the 'execution ordering' of the opcodes
@@ -85,16 +90,15 @@ impl MergeExpressionsOptimizer {
85
90
// - doing this pass again until there is no change, or
86
91
// - merging 'b' into 'i' instead
87
92
if i < b {
88
- if let Some ( expr) = Self :: merge ( & expr_use, & expr_define, w) {
93
+ if let Some ( expr) = Self :: merge ( expr_use, expr_define, w) {
89
94
modified_gates. insert ( b, Opcode :: AssertZero ( expr) ) ;
90
95
to_keep = false ;
91
96
// Update the 'used_witness' map to account for the merge.
92
- for w2 in CircuitSimulator :: expr_wit ( & expr_define) {
97
+ for w2 in CircuitSimulator :: expr_wit ( expr_define) {
93
98
if !circuit_inputs. contains ( & w2) {
94
- let mut v = used_witness[ & w2 ] . clone ( ) ;
99
+ let v = used_witness. entry ( w2 ) . or_default ( ) ;
95
100
v. insert ( b) ;
96
101
v. remove ( & i) ;
97
- used_witness. insert ( w2, v) ;
98
102
}
99
103
}
100
104
// We need to stop here and continue with the next opcode
@@ -107,12 +111,9 @@ impl MergeExpressionsOptimizer {
107
111
}
108
112
109
113
if to_keep {
110
- if modified_gates. contains_key ( & i) {
111
- new_circuit. push ( modified_gates[ & i] . clone ( ) ) ;
112
- } else {
113
- new_circuit. push ( opcode. clone ( ) ) ;
114
- }
115
- new_acir_opcode_positions. push ( acir_opcode_positions[ i] ) ;
114
+ let opcode = modified_gates. get ( & i) . cloned ( ) . unwrap_or ( opcode) ;
115
+ new_circuit. push ( opcode) ;
116
+ new_acir_opcode_positions. push ( opcode_position) ;
116
117
}
117
118
}
118
119
( new_circuit, new_acir_opcode_positions)
0 commit comments