1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
/*++
Copyright (c) 2026 Romain Soulat
Module Name:
ff_rewriter.h
Abstract:
Rewriter for prime-field arithmetic, shared by simplification and model
evaluation. Normalizes sums and products over F_p without distributing
products of symbolic sums or cancelling unknown factors.
Author:
Romain Soulat
// Shared by simplification and model evaluation. This normalizes sums and
// products without distributing products of symbolic sums or cancelling unknown factors.
// All identities are in F_p. Inverse-based rules require p prime and a nonzero
// constant coefficient; accepting a large modulus is not a primality certificate.
--*/
#pragma once
#include "ast/ff_decl_plugin.h"
#include "ast/occurs.h"
#include "ast/rewriter/rewriter_types.h"
#include "util/common_msgs.h"
#include "util/z3_exception.h"
#include <map>
#include <vector>
#include <algorithm>
class ff_rewriter {
ast_manager &m;
ff_util u;
struct expr_order {
bool operator()(expr *a, expr *b) const {
return a->get_id() < b->get_id();
}
};
struct monomial_order {
ff_util const &u;
bool operator()(expr *a, expr *b) const {
bool am = u.is_mul(a), bm = u.is_mul(b);
unsigned na = am ? to_app(a)->get_num_args() : 1;
unsigned nb = bm ? to_app(b)->get_num_args() : 1;
if (na != nb)
return na < nb;
for (unsigned i = 0; i < na; ++i) {
expr *x = am ? to_app(a)->get_arg(i) : a;
expr *y = bm ? to_app(b)->get_arg(i) : b;
if (x != y)
return x->get_id() < y->get_id();
}
return false;
}
};
// A coefficient-free product can be temporary. Order by its retained
// factors, not by the temporary product's recyclable AST id.
using coefficients = std::map<expr *, rational, monomial_order>;
expr_ref product(expr_ref_vector const &factors, rational const &c, sort *s) {
// 0*t = 0 and 1*t = t. The empty product is 1, so a product with no
// symbolic factors denotes c itself; a singleton needs no MUL node.
expr_ref_vector args(m);
if (c.is_zero())
return expr_ref(u.mk_numeral(c, s), m);
if (!c.is_one() || factors.empty())
args.push_back(u.mk_numeral(c, s));
args.append(factors);
if (args.size() == 1)
return expr_ref(args.get(0), m);
return expr_ref(u.mk_mul(args), m);
}
// Children have already been rewritten. Temporary products must remain
// referenced while their raw pointers serve as coefficient-map keys.
// Associativity and distributivity give k*(a+b) = k*a+k*b. Extract only
// numerical scalars: a product of symbolic sums remains an atomic term.
// Reducing coefficients modulo p preserves their value in F_p.
void collect(expr *e, rational scale, rational const &p, coefficients &terms, rational &constant,
expr_ref_vector &pins, unsigned &count) {
std::vector<std::pair<expr *, rational>> todo;
todo.emplace_back(e, scale);
while (!todo.empty()) {
if (!m.inc())
throw rewriter_exception(Z3_CANCELED_MSG);
auto [a, factor] = todo.back();
todo.pop_back();
if (u.is_add(a)) {
for (expr *arg : *to_app(a))
todo.emplace_back(arg, factor);
continue;
}
rational value;
if (u.is_numeral(a, value)) {
constant = mod(constant + factor * value, p);
continue;
}
rational coeff = factor;
expr_ref base(a, m);
if (u.is_mul(a)) {
expr_ref_vector factors(m);
for (expr *arg : *to_app(a)) {
if (u.is_numeral(arg, value))
coeff = mod(coeff * value, p);
else
factors.push_back(arg);
}
base = product(factors, rational(1), a->get_sort());
pins.push_back(base);
}
if (u.is_add(base)) {
for (expr *arg : *to_app(base))
todo.emplace_back(arg, coeff);
continue;
}
if (u.is_numeral(base, value))
constant = mod(constant + coeff * value, p);
else {
++count;
// c*t+d*t = (c+d)*t, including t-t = 0. This is additive
// cancellation and does not assume that t is nonzero.
auto [it, fresh] = terms.try_emplace(base.get(), rational(0));
it->second = mod(it->second + coeff, p);
if (it->second.is_zero())
terms.erase(it);
}
}
}
expr_ref sum(coefficients const &terms, rational const &constant, sort *s) {
// Addition is commutative, 0+t = t, and an empty sum is 0. Reordering
// summands and omitting zero coefficients therefore preserve the term.
expr_ref_vector args(m);
if (!constant.is_zero())
args.push_back(u.mk_numeral(constant, s));
for (auto const &[term, coeff] : terms) {
expr_ref_vector factors(m);
// Flatten a monomial when restoring its coefficient.
if (u.is_mul(term))
for (expr *arg : *to_app(term))
factors.push_back(arg);
else
factors.push_back(term);
args.push_back(product(factors, coeff, s));
}
if (args.empty())
return expr_ref(u.mk_numeral(rational(0), s), m);
if (args.size() == 1)
return expr_ref(args.get(0), m);
return expr_ref(u.mk_add(args), m);
}
rational inverse(rational const &a, rational const &p) {
// 1 and -1 are self-inverse (also in F_2). Otherwise extended Euclid
// yields t*a + k*p = 1, hence t*a = 1 in F_p for nonzero a.
if (a.is_one() || a == p - rational(1))
return a;
rational t, k, g = gcd(a, p, t, k);
// The sort contract requires p prime; the caller supplies nonzero a.
SASSERT(g.is_one());
return mod(t, p);
}
public:
explicit ff_rewriter(ast_manager &m) : m(m), u(m) {}
family_id get_fid() const {
return u.get_fid();
}
br_status mk_eq_core(expr *a, expr *b, expr_ref &out) {
rational const &p = u.modulus(a->get_sort());
coefficients terms(monomial_order{u});
rational constant(0);
expr_ref_vector pins(m);
unsigned count = 0;
expr_ref eq(m.mk_eq(a, b), m);
// a=b iff a-b=0. If all symbolic summands cancel, this is exactly
// the test that the remaining canonical constant is zero.
collect(a, rational(1), p, terms, constant, pins, count);
collect(b, rational(-1), p, terms, constant, pins, count);
if (terms.empty()) {
out = m.mk_bool_val(constant.is_zero());
return out == eq ? BR_FAILED : BR_REWRITE_FULL;
}
if (terms.size() == 1) {
// c*t+k=0 iff t=-k/c because c is a nonzero field constant.
// t may itself be a product or sum; no factor of t is cancelled.
auto const &[base, coeff] = *terms.begin();
expr_ref value(u.mk_numeral(mod(-constant * inverse(coeff, p), p), a->get_sort()), m);
out = m.mk_eq(base, value);
return out == eq ? BR_FAILED : BR_REWRITE_FULL;
}
// Expose affine wire definitions to solve-eqs even when a circuit
// exporter writes t = c - a*x. Prefer a wire absent from the other
// terms: orienting x = x*y + z would hide the usable definition of z.
// For a nonzero constant c, c*x+r=0 iff x=-c^{-1}*r. The occurs
// check makes this a nonrecursive definition usable for substitution.
// Preferring an existing left side or a unit coefficient is only a
// choice of orientation, not an additional algebraic premise.
auto pivot = terms.end();
bool pivot_unit = false;
expr *preferred = nullptr;
if (is_uninterp_const(a) != is_uninterp_const(b)) {
expr *v = is_uninterp_const(a) ? a : b;
expr *rhs = v == a ? b : a;
if (!occurs(v, rhs))
preferred = v;
}
for (auto it = terms.begin(); it != terms.end(); ++it) {
if (!is_uninterp_const(it->first))
continue;
if (preferred && it->first != preferred)
continue;
bool unit = it->second.is_one() || it->second == p - rational(1);
if (pivot_unit && !unit)
continue;
bool safe = true;
for (auto const &[term, coeff] : terms)
if (term != it->first && occurs(it->first, term)) {
safe = false;
break;
}
if (safe) {
pivot = it;
pivot_unit = unit;
}
}
if (pivot != terms.end()) {
expr_ref var(pivot->first, m);
rational factor = mod(-inverse(pivot->second, p), p);
terms.erase(pivot);
for (auto &[term, coeff] : terms)
coeff = mod(coeff * factor, p);
expr_ref rhs = sum(terms, mod(constant * factor, p), a->get_sort());
out = m.mk_eq(var, rhs);
return out == eq ? BR_FAILED : BR_REWRITE_FULL;
}
// Avoid expanding an ordinary x=y or changing equalities that contain
// no cancellable summands. Never cancel a symbolic multiplicative factor.
// The residual equation sum(c_i*t_i)=-k is the same a-b=0 identity.
// In particular, x*y=x*z cannot become y=z without knowing x != 0.
if (terms.size() >= count)
return BR_FAILED;
expr_ref left = sum(terms, rational(0), a->get_sort());
expr_ref right(u.mk_numeral(mod(-constant, p), a->get_sort()), m);
out = m.mk_eq(left, right);
return out == eq ? BR_FAILED : BR_REWRITE_FULL;
}
br_status mk_app_core(func_decl *f, unsigned n, expr *const *args, expr_ref &out) {
auto kind = f->get_decl_kind();
if (kind == OP_FF_NUM)
return BR_FAILED;
sort *s = f->get_range();
rational const &p = u.modulus(s);
if (kind == OP_FF_ADD) {
coefficients terms(monomial_order{u});
rational constant(0);
expr_ref_vector pins(m);
unsigned count = 0;
for (unsigned i = 0; i < n; ++i)
collect(args[i], rational(1), p, terms, constant, pins, count);
out = sum(terms, constant, s);
return BR_DONE;
}
if (kind == OP_FF_MUL || kind == OP_FF_NEG) {
// -t=(-1)*t; associativity flattens products and commutativity
// permits sorting their factors. Fold constants modulo p, stop
// at the annihilator 0, and use (-1)*(-1)=1 for double negation.
rational coeff(kind == OP_FF_NEG ? -1 : 1), value;
expr_ref_vector factors(m);
ptr_vector<expr> todo;
for (unsigned i = 0; i < n; ++i)
todo.push_back(args[i]);
while (!todo.empty()) {
if (!m.inc())
throw rewriter_exception(Z3_CANCELED_MSG);
expr *a = todo.back();
todo.pop_back();
if (u.is_numeral(a, value))
coeff = mod(coeff * value, p);
else if (u.is_mul(a))
for (expr *arg : *to_app(a))
todo.push_back(arg);
else
factors.push_back(a);
if (coeff.is_zero())
break;
}
ptr_vector<expr> ordered;
for (expr *e : factors)
ordered.push_back(e);
std::stable_sort(ordered.begin(), ordered.end(), expr_order());
expr_ref_vector sorted(m);
for (expr *e : ordered)
sorted.push_back(e);
out = product(sorted, mod(coeff, p), s);
return BR_DONE;
}
if (kind == OP_FF_BITSUM) {
// Preserve bit positions: removing a zero operand would change weights.
// By definition bitsum(a_0,...,a_n)=sum(2^i*a_i) in F_p.
// Constant evaluation needs neither Boolean arguments nor a no-wrap
// bound; it does not infer either condition for symbolic arguments.
rational value, total(0), weight(1);
for (unsigned i = 0; i < n; ++i) {
if (!u.is_numeral(args[i], value))
return BR_FAILED;
total = mod(total + weight * value, p);
weight = mod(weight * rational(2), p);
}
out = u.mk_numeral(total, s);
return BR_DONE;
}
return BR_FAILED;
}
};