diff --git a/.gitignore b/.gitignore index b889815..e550925 100644 --- a/.gitignore +++ b/.gitignore @@ -163,3 +163,7 @@ cython_debug/ ignore m4ri-release-* +# setup.py downloads m4ri-{release}/ on build +m4ri-*/ +# local uv lockfile +uv.lock diff --git a/benchmarks/mt_benchmark.py b/benchmarks/mt_benchmark.py new file mode 100644 index 0000000..b3e0555 --- /dev/null +++ b/benchmarks/mt_benchmark.py @@ -0,0 +1,91 @@ +import argparse +import json +import random +import statistics +from time import perf_counter + +from gf2bv import LinearSystem +from gf2bv.crypto.mt import MT19937 + + +def run_mt_case(bs: int, samples: int | None = None) -> dict[str, float | bool | int]: + rand = random.Random(3142) + state = tuple(rand.getstate()[1][:-1]) + + effective_bs = ((bs - 1) & bs) or bs + samples = 624 * 32 // effective_bs if samples is None else samples + outputs = [rand.getrandbits(bs) for _ in range(samples)] + + lin = LinearSystem([32] * 624) + mt = lin.gens() + rng = MT19937(mt) + + build_start = perf_counter() + zeros = [rng.getrandbits(bs) ^ out for out in outputs] + [mt[0] ^ 0x80000000] + build_end = perf_counter() + solution = lin.solve_one(zeros) + solve_end = perf_counter() + + return { + "build_seconds": build_end - build_start, + "solve_seconds": solve_end - build_end, + "total_seconds": solve_end - build_start, + "samples": samples, + "ok": solution == state, + } + + +CASES: dict[str, tuple[int, int | None]] = { + "mt32": (32, None), + "mt17": (17, None), + "mt9": (9, None), + "mt1": (1, None), + "mt1337": (1337, 19968 // 1337 + 10), + "mt137": (137, 19968 // 137 + 60), +} + + +def summarize(case_name: str, runs: list[dict[str, float | bool | int]]) -> dict[str, object]: + if not all(run["ok"] for run in runs): + raise AssertionError(f"{case_name} produced an incorrect solution") + + build = [float(run["build_seconds"]) for run in runs] + solve = [float(run["solve_seconds"]) for run in runs] + total = [float(run["total_seconds"]) for run in runs] + + return { + "runs": runs, + "summary": { + "build_mean": statistics.fmean(build), + "build_median": statistics.median(build), + "solve_mean": statistics.fmean(solve), + "solve_median": statistics.median(solve), + "total_mean": statistics.fmean(total), + "total_median": statistics.median(total), + }, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--runs", type=int, default=10) + parser.add_argument( + "--case", + dest="cases", + action="append", + choices=sorted(CASES), + help="Benchmark only the selected case. Repeat the flag to include multiple cases.", + ) + args = parser.parse_args() + + selected_cases = args.cases or list(CASES) + results: dict[str, object] = {"runs": args.runs, "cases": {}} + for case_name in selected_cases: + bs, samples = CASES[case_name] + case_runs = [run_mt_case(bs, samples) for _ in range(args.runs)] + results["cases"][case_name] = summarize(case_name, case_runs) + print(json.dumps(results, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/gf2bv/__init__.py b/gf2bv/__init__.py index a3c3fa6..2d3fd80 100644 --- a/gf2bv/__init__.py +++ b/gf2bv/__init__.py @@ -232,11 +232,9 @@ def _solve_internal(self, zeros: Zeros, mode: TSolveMode): # no solution return cols = self._cols - if cols > len(eqs): - # m4ri_solve requires rows >= cols, pad with zeros - eqs += [0] * (cols - len(eqs)) # all mode (mode == 1): may return None if no solution, otherwise return an iterator # one solution mode (mode == 0): return the solution directly if it exists, otherwise return None + # the native solver pads missing rows with zeros internally return m4ri_solve(eqs, cols, mode) def _convert_sol(self, s: int) -> tuple[int, ...]: diff --git a/gf2bv/_internal.c b/gf2bv/_internal.c index 32799c4..560b642 100644 --- a/gf2bv/_internal.c +++ b/gf2bv/_internal.c @@ -9,10 +9,12 @@ #define NON_SIZE_BITS 3 #define PyLong_DigitCount(o) ((o)->long_value.lv_tag >> NON_SIZE_BITS) #define GET_OB_DIGITS(o) ((o)->long_value.ob_digit) +#define PyLong_IsNegative(o) (((o)->long_value.lv_tag & _PyLong_SIGN_MASK) == 2) #else // pre 3.12 #define PyLong_DigitCount(o) Py_ABS(Py_SIZE(o)) #define GET_OB_DIGITS(o) ((o)->ob_digit) +#define PyLong_IsNegative(o) (Py_SIZE(o) < 0) #endif #if PY_VERSION_HEX >= 0x030C0000 @@ -38,6 +40,81 @@ static inline PyObject *mzd_vector_to_pylong(char *buf, mzd_t *v) { return PyLong_FromString(buf, NULL, 2); } +static inline word low_bitmask(int bits) { + if (bits <= 0) + return 0; + if (bits >= 64) + return UINT64_MAX; + return (((word)1) << bits) - 1; +} + +static inline BIT pylong_pack_mzd_row_skip_lsb(PyLongObject *value, + word *row, + wi_t width) { + Py_ssize_t n_digits = PyLong_DigitCount(value); + if (n_digits == 0) { + for (wi_t word_idx = 0; word_idx < width; word_idx++) { + row[word_idx] = 0; + } + return 0; + } + + const digit *digits = GET_OB_DIGITS(value); + BIT low = digits[0] & 1; + word spill = ((word)digits[0]) >> 1; + int spill_bits = PyLong_SHIFT - 1; + Py_ssize_t digit_idx = 1; + + for (wi_t word_idx = 0; word_idx < width; word_idx++) { + word packed = 0; + int filled = 0; + + if (spill_bits > 0) { + int take = spill_bits < 64 ? spill_bits : 64; + packed = spill & low_bitmask(take); + spill >>= take; + spill_bits -= take; + filled = take; + } + + while (filled < 64 && digit_idx < n_digits) { + word next = digits[digit_idx++]; + if (filled <= 64 - PyLong_SHIFT) { + packed |= next << filled; + filled += PyLong_SHIFT; + } else { + int take = 64 - filled; + packed |= (next & low_bitmask(take)) << filled; + spill = next >> take; + spill_bits = PyLong_SHIFT - take; + filled = 64; + } + } + + row[word_idx] = packed; + } + + return low; +} + +static inline int object_bit_constant(PyObject *obj) { + if (obj == Py_False) + return 0; + if (obj == Py_True) + return 1; + if (!PyLong_Check(obj)) + return -1; + PyLongObject *value = (PyLongObject *)obj; + if (PyLong_IsNegative(value)) + return -1; + Py_ssize_t n_digits = PyLong_DigitCount(value); + if (n_digits == 0) + return 0; + if (n_digits == 1 && GET_OB_DIGITS(value)[0] == 1) + return 1; + return -1; +} + #define Iter_PyLong_Bits(obj, max, action, after_action) \ do { \ Py_ssize_t _n_digits = PyLong_DigitCount(obj); \ @@ -356,6 +433,94 @@ mzd_t *_mzd_kernel_left_pluq(mzd_t *A, return R; } +// Fast path for SOLVE_MODE_SINGLE: build [A | b] and RREF via M4RI. +// Returns a Python int solution or Py_None if inconsistent. +static PyObject *m4ri_solve_single_augmented(PyObject *linsys_list, + Py_ssize_t cols) { + Py_ssize_t used_rows = PyList_GET_SIZE(linsys_list); + Py_ssize_t rows = used_rows > cols ? used_rows : cols; + + // Aug = [A | b] where the last column holds the affine (b) bit. + mzd_t *Aug = mzd_init(rows, cols + 1); + if (Aug == NULL) { + return PyErr_NoMemory(); + } + wi_t pack_width = (cols + 63) / 64; + int b_col_word = cols / 64; + int b_col_bit = cols % 64; + // Mask to clear any stray bits in A's last packed word above column cols-1. + // pylong_pack_mzd_row_skip_lsb fills pack_width*64 bits, which may exceed + // cols when cols is not a multiple of 64, so a malformed input with extra + // high bits could otherwise bleed into the b column. + word a_high_mask = low_bitmask(cols % 64 ? cols % 64 : 64); + + for (Py_ssize_t r = 0; r < used_rows; r++) { + PyObject *item = PyList_GET_ITEM(linsys_list, r); + if (!PyLong_Check(item)) { + mzd_free(Aug); + PyErr_SetString(PyExc_TypeError, "List items must be integers"); + return NULL; + } + PyLongObject *v = (PyLongObject *)item; + if (PyLong_DigitCount(v) == 0) + continue; + word *row = mzd_row(Aug, r); + BIT low = pylong_pack_mzd_row_skip_lsb(v, row, pack_width); + row[pack_width - 1] &= a_high_mask; + if (low) + row[b_col_word] |= ((word)1) << b_col_bit; + } + + PyThreadState *_save = PyEval_SaveThread(); + rci_t rank = mzd_echelonize_m4ri(Aug, 1, 0); + + // After RREF, each of the first `rank` rows has a unique leading 1. If + // any leading 1 is in the last (augmented b) column, the system is + // inconsistent. Otherwise variable j's value equals the b-column bit of + // the row whose pivot is column j. + mzd_t *sol0 = mzd_init(1, cols); + if (sol0 == NULL) { + mzd_free(Aug); + PyEval_RestoreThread(_save); + return PyErr_NoMemory(); + } + int inconsistent = 0; + for (rci_t i = 0; i < rank; i++) { + word *row = mzd_row(Aug, i); + rci_t pivot = -1; + for (wi_t w = 0; w < Aug->width; w++) { + if (row[w]) { + pivot = w * 64 + __builtin_ctzll(row[w]); + break; + } + } + if (pivot == cols) { + inconsistent = 1; + break; + } + if ((row[b_col_word] >> b_col_bit) & 1) + mzd_write_bit(sol0, 0, pivot, 1); + } + mzd_free(Aug); + PyEval_RestoreThread(_save); + + if (inconsistent) { + mzd_free(sol0); + return PythonNone; + } + + char *str = malloc(cols + 1); + if (str == NULL) { + mzd_free(sol0); + return PyErr_NoMemory(); + } + str[cols] = '\0'; + PyObject *ret = mzd_vector_to_pylong(str, sol0); + free(str); + mzd_free(sol0); + return ret; +} + PyObject *m4ri_solve(PyObject *self, PyObject *const *args, Py_ssize_t nargs) { PyObject *linsys_list; Py_ssize_t cols; @@ -386,13 +551,11 @@ PyObject *m4ri_solve(PyObject *self, PyObject *const *args, Py_ssize_t nargs) { PyErr_SetString(PyExc_ValueError, "Invalid mode"); return NULL; } - Py_ssize_t rows = PyList_GET_SIZE(linsys_list); - if (rows < cols) { - PyErr_SetString(PyExc_ValueError, - "Number of rows must be greater than or equal to " - "number of columns, try pad with zeros."); - return NULL; + if (mode == SOLVE_MODE_SINGLE) { + return m4ri_solve_single_augmented(linsys_list, cols); } + Py_ssize_t used_rows = PyList_GET_SIZE(linsys_list); + Py_ssize_t rows = used_rows > cols ? used_rows : cols; // Ax = B mzd_t *A = mzd_init(rows, cols); @@ -400,7 +563,7 @@ PyObject *m4ri_solve(PyObject *self, PyObject *const *args, Py_ssize_t nargs) { int B_is_not_zero = 0; - for (Py_ssize_t r = 0; r < rows; r++) { + for (Py_ssize_t r = 0; r < used_rows; r++) { PyObject *item = PyList_GET_ITEM(linsys_list, r); if (!PyLong_Check(item)) { mzd_free(A); @@ -413,16 +576,14 @@ PyObject *m4ri_solve(PyObject *self, PyObject *const *args, Py_ssize_t nargs) { // row in A: [v_1 v_2 ... v_cols] // row in B: [v_0] PyLongObject *v = (PyLongObject *)item; - Iter_PyLong_Bits(v, cols + 1, - { - if (bitcnt == 0) { - mzd_write_bit(B, r, 0, bit); - B_is_not_zero |= bit; - } else { - mzd_write_bit(A, r, bitcnt - 1, bit); - } - }, - {}); + if (PyLong_DigitCount(v) == 0) { + continue; + } + word *A_row = mzd_row(A, r); + BIT low = pylong_pack_mzd_row_skip_lsb(v, A_row, A->width); + A_row[A->width - 1] &= A->high_bitmask; + mzd_row(B, r)[0] = low; + B_is_not_zero |= low; } // release GIL because we don't need to call Python API now @@ -603,6 +764,92 @@ PyObject *mul_bit_quad(PyObject *self, return v; } +// Fused: result[i] = bits[i] XOR ((shifted)[i] AND mask_bit[i]) +// where (shifted)[i] = bits[i+n] (right shift) or bits[i-n] (left shift), +// clamped to 0 when the index is out of range. +// This replaces the common pattern `y ^= (y >> n) & mask` (or `<<`) with a +// single pass, avoiding two intermediate tuple allocations and their +// dispatch overhead. mask must be a non-negative Python int; its bits 0..N-1 +// are consulted, where N = len(bits). +PyObject *bv_xor_shift_mask(PyObject *self, + PyObject *const *args, + Py_ssize_t nargs) { + if (nargs != 4) { + PyErr_SetString(PyExc_TypeError, + "bv_xor_shift_mask requires 4 arguments"); + return NULL; + } + PyObject *bits = args[0]; + if (!PyTuple_Check(bits)) { + PyErr_SetString(PyExc_TypeError, "bits must be a tuple"); + return NULL; + } + Py_ssize_t shift = PyLong_AsSsize_t(args[1]); + if (shift < 0) { + if (shift == -1 && PyErr_Occurred()) + return NULL; + PyErr_SetString(PyExc_ValueError, "shift must be non-negative"); + return NULL; + } + int is_right = PyObject_IsTrue(args[2]); + if (is_right < 0) + return NULL; + PyObject *mask_obj = args[3]; + if (!PyLong_Check(mask_obj)) { + PyErr_SetString(PyExc_TypeError, "mask must be an integer"); + return NULL; + } + PyLongObject *mask_long = (PyLongObject *)mask_obj; + if (PyLong_IsNegative(mask_long)) { + PyErr_SetString(PyExc_ValueError, "mask must be non-negative"); + return NULL; + } + Py_ssize_t mask_ndig = PyLong_DigitCount(mask_long); + const digit *mask_digits = GET_OB_DIGITS(mask_long); + + Py_ssize_t n = PyTuple_GET_SIZE(bits); + PyObject *ret = PyTuple_New(n); + if (ret == NULL) + return NULL; + for (Py_ssize_t i = 0; i < n; i++) { + Py_ssize_t digit_idx = i / PyLong_SHIFT; + int mask_bit = 0; + if (digit_idx < mask_ndig) { + int bit_in_digit = (int)(i % PyLong_SHIFT); + mask_bit = (int)((mask_digits[digit_idx] >> bit_in_digit) & 1); + } + PyObject *a_item = PyTuple_GET_ITEM(bits, i); + PyObject *result_item; + Py_ssize_t src = is_right ? i + shift : i - shift; + if (mask_bit == 0 || src < 0 || src >= n) { + result_item = Py_NewRef(a_item); + } else { + PyObject *b_item = PyTuple_GET_ITEM(bits, src); + int bit_a = object_bit_constant(a_item); + int bit_b = object_bit_constant(b_item); + if (bit_a == 0) { + result_item = Py_NewRef(b_item); + } else if (bit_b == 0) { + result_item = Py_NewRef(a_item); + } else if (bit_a == 1 && bit_b == 1) { + result_item = PythonFalse; + } else if (bit_a == 1) { + result_item = PyNumber_Xor(b_item, Py_True); + } else if (bit_b == 1) { + result_item = PyNumber_Xor(a_item, Py_True); + } else { + result_item = PyNumber_Xor(a_item, b_item); + } + if (result_item == NULL) { + Py_DECREF(ret); + return NULL; + } + } + PyTuple_SET_ITEM(ret, i, result_item); + } + return ret; +} + PyObject *xor_tuple(PyObject *self, PyObject *const *args, Py_ssize_t nargs) { PyObject *a, *b; if (nargs != 2) { @@ -625,7 +872,22 @@ PyObject *xor_tuple(PyObject *self, PyObject *const *args, Py_ssize_t nargs) { for (Py_ssize_t i = 0; i < len_a; i++) { PyObject *item_a = PyTuple_GET_ITEM(a, i); PyObject *item_b = PyTuple_GET_ITEM(b, i); - PyObject *xor_item = PyNumber_Xor(item_a, item_b); + int bit_a = object_bit_constant(item_a); + int bit_b = object_bit_constant(item_b); + PyObject *xor_item; + if (bit_a == 0) { + xor_item = Py_NewRef(item_b); + } else if (bit_b == 0) { + xor_item = Py_NewRef(item_a); + } else if (bit_a == 1 && bit_b == 1) { + xor_item = PythonFalse; + } else if (bit_a == 1) { + xor_item = PyNumber_Xor(item_b, Py_True); + } else if (bit_b == 1) { + xor_item = PyNumber_Xor(item_a, Py_True); + } else { + xor_item = PyNumber_Xor(item_a, item_b); + } if (xor_item == NULL) { Py_DECREF(ret); PyErr_SetString( @@ -781,6 +1043,15 @@ static PyMethodDef methods[] = { "\n" "Multiply two linear symbolic bits to a linearized quadratic symbolic " "bit"}, + {"bv_xor_shift_mask", _PyCFunction_CAST(bv_xor_shift_mask), METH_FASTCALL, + "bv_xor_shift_mask(bits, shift, is_right, mask)\n" + "--\n" + "\n" + "Fused `bits ^ ((bits shifted) & mask)`: returns a new tuple. For " + "is_right=True the shift is right (index i pulls from i+shift); for " + "is_right=False it is left (index i pulls from i-shift). Out-of-range " + "source indices contribute zero. mask must be a non-negative int; only " + "its low N bits are consulted where N = len(bits)."}, {"xor_tuple", _PyCFunction_CAST(xor_tuple), METH_FASTCALL, "xor_tuple(a, b)\n" "--\n" diff --git a/gf2bv/crypto/mt.py b/gf2bv/crypto/mt.py index 7be4bf3..2eabaa3 100644 --- a/gf2bv/crypto/mt.py +++ b/gf2bv/crypto/mt.py @@ -1,6 +1,7 @@ import random from .. import BitVec +from .._internal import bv_xor_shift_mask class MersenneTwister: @@ -39,6 +40,16 @@ def twist(self): self.mt[i] = self.mt[(i + self.m) % self.n] ^ (y >> 1) ^ sel def temper(self, y): + if isinstance(y, BitVec) and len(y._bits) == self.w: + # Fused path: each `y ^= (y shift) & mask` becomes one C call, + # avoiding two intermediate tuple allocations per stage. This is + # the hot path for symbolic MT state recovery. + bits = y._bits + bits = bv_xor_shift_mask(bits, self.u, True, self.d) + bits = bv_xor_shift_mask(bits, self.s, False, self.w1 & self.b) + bits = bv_xor_shift_mask(bits, self.t, False, self.w1 & self.c) + bits = bv_xor_shift_mask(bits, self.l, True, self.w1) + return BitVec(bits) y ^= (y >> self.u) & self.d y ^= (y << self.s) & self.w1 & self.b y ^= (y << self.t) & self.w1 & self.c @@ -53,8 +64,59 @@ def __call__(self): self.mti += 1 return self.temper(y) + @property + def _temper_output_map(self): + # Lazy, per-instance cache. For each output bit index p, list the input + # bit indices whose XOR produces temper(y) bit p. Computed by running + # the integer temper on each basis vector 1<> self.u) & self.d + y ^= (y << self.s) & self.w1 & self.b + y ^= (y << self.t) & self.w1 & self.c + y ^= y >> self.l + for p in range(self.w): + if (y >> p) & 1: + maps[p].append(i) + self.__dict__["_temper_output_map_cache"] = maps + return maps + def _getrandbits_word(self, k): - r = self() + # Fast path for symbolic BitVec when few output bits are requested: + # compute only those bits via the precomputed temper linear map. For + # MT19937 the average hamming weight per output bit is ~5, so for + # small k this beats the full-temper path which always does ~59 + # PyLong XORs regardless of k. The threshold of half the word + # width keeps the break-even on the cautious side — beyond that the + # full temper reuses intermediate results and wins. + if self.mti >= self.n: + self.twist() + self.mti = 0 + y = self.mt[self.mti] + self.mti += 1 + if ( + isinstance(y, BitVec) + and len(y._bits) == self.w + and k <= self.w // 2 + ): + maps = self._temper_output_map + bits = y._bits + out: list[int] = [] + for p in range(self.w - k, self.w): + indices = maps[p] + if not indices: + out.append(0) + continue + r = bits[indices[0]] + for idx in indices[1:]: + r = r ^ bits[idx] + out.append(r) + return BitVec(tuple(out)) + r = self.temper(y) if isinstance(r, BitVec): return r[self.w - k :] return r >> (self.w - k) diff --git a/setup.py b/setup.py index 84e8037..542777a 100644 --- a/setup.py +++ b/setup.py @@ -31,22 +31,55 @@ def download_and_build_m4ri(): if not workdir.exists(): raise FileNotFoundError(f"Failed to extract {workdir}") subprocess.run("autoreconf --install", shell=True, cwd=workdir, check=True) - subprocess.run( - './configure CFLAGS="-fPIC -O3 -march=native -mtune=native" --enable-openmp --enable-thread-safe', - shell=True, - cwd=workdir, - check=True, - ) + if sys.platform.startswith("darwin"): + # Apple clang's OpenMP needs `-Xpreprocessor -fopenmp`, which + # m4ri's configure doesn't probe for. Configure without its + # OpenMP detection and patch m4ri_config.h + CFLAGS manually. + libomp_prefix = os.environ.get( + "LIBOMP_PREFIX", "/opt/homebrew/opt/libomp" + ) + omp_cflags = f"-Xpreprocessor -fopenmp -I{libomp_prefix}/include" + omp_ldflags = f"-L{libomp_prefix}/lib -lomp" + subprocess.run( + f'./configure CFLAGS="-fPIC -O3 -march=native -mtune=native {omp_cflags}" ' + f'LDFLAGS="{omp_ldflags}" --enable-thread-safe', + shell=True, + cwd=workdir, + check=True, + ) + # Flip __M4RI_HAVE_OPENMP=1 so the OMP pragmas compile in; + # m4ri's autoconf probe for OpenMP fails on Apple clang because + # it only tries bare `-fopenmp`. + config_h = workdir / "m4ri" / "m4ri_config.h" + text = config_h.read_text() + config_h.write_text( + text.replace( + "#define __M4RI_HAVE_OPENMP\t\t0", + "#define __M4RI_HAVE_OPENMP\t\t1", + ) + ) + else: + subprocess.run( + './configure CFLAGS="-fPIC -O3 -march=native -mtune=native" --enable-openmp --enable-thread-safe', + shell=True, + cwd=workdir, + check=True, + ) subprocess.run("make -j", shell=True, cwd=workdir, check=True) if not libm4ri_a.exists(): raise FileNotFoundError(f"Failed to build {libm4ri_a}") extra_compile_args = ["-O3", "-march=native", "-mtune=native"] extra_link_args = [] if sys.platform.startswith("darwin"): - # you may want to use the following line if you installed libomp via Homebrew - # export LDFLAGS="-L/opt/homebrew/opt/libomp/lib" - extra_compile_args += ["-Xpreprocessor", "-fopenmp"] - extra_link_args += ["-lomp"] + libomp_prefix = os.environ.get( + "LIBOMP_PREFIX", "/opt/homebrew/opt/libomp" + ) + extra_compile_args += [ + "-Xpreprocessor", + "-fopenmp", + f"-I{libomp_prefix}/include", + ] + extra_link_args += [f"-L{libomp_prefix}/lib", "-lomp"] elif not sys.platform.startswith("win"): extra_compile_args += ["-fopenmp"] extra_link_args += ["-fopenmp"] diff --git a/tests/test_solver.py b/tests/test_solver.py new file mode 100644 index 0000000..e03e0d2 --- /dev/null +++ b/tests/test_solver.py @@ -0,0 +1,185 @@ +"""Regression tests for the two m4ri_solve paths (single vs affine-space). + +Covers the three shapes that the reviewer flagged as lacking coverage: + +1. Satisfiable, over-determined — the common benchmark shape. +2. Inconsistent — must return None in both modes. +3. Underdetermined — the fast `solve_one` path must return *some* + valid solution (not necessarily unique), and + the `solve_raw_space` path must describe the + full affine space containing it. +""" + +from __future__ import annotations + +import random + +import pytest + +from gf2bv import LinearSystem +from gf2bv._internal import m4ri_solve + + +def _verify(eqs: list[int], cols: int, sol: int) -> bool: + """Check sol satisfies every equation: bit0(e) ⊕ popcount((e>>1) & sol) == 0.""" + for e in eqs: + affine = e & 1 + linear = (e >> 1) & ((1 << cols) - 1) + if affine ^ ((linear & sol).bit_count() & 1) != 0: + return False + return True + + +def _in_span(target: int, basis) -> bool: + """True iff `target` is in the GF(2)-linear span of `basis`. + + Incremental RREF: for each basis vector, reduce it by the accumulated + pivot set, then if non-zero record it as a new pivot at its leading bit. + Finally reduce `target` by the pivot set — it lies in the span iff it + reduces to 0. + """ + pivots: list[int] = [] + leads: list[int] = [] + for v in basis: + v = int(v) + for lead, p in zip(leads, pivots): + if (v >> lead) & 1: + v ^= p + if v == 0: + continue + pivots.append(v) + leads.append(v.bit_length() - 1) + t = int(target) + for lead, p in zip(leads, pivots): + if (t >> lead) & 1: + t ^= p + return t == 0 + + +def _random_system( + n_vars: int, + n_eqs: int, + *, + seed: int, + planted: bool = True, +) -> tuple[list[int], int | None]: + """Build `n_eqs` random GF(2) equations on `n_vars` variables. + + When `planted`, a random solution is generated and every equation is made + consistent with it, so the system is always satisfiable. + """ + rng = random.Random(seed) + sol = rng.getrandbits(n_vars) if planted else None + eqs: list[int] = [] + for _ in range(n_eqs): + coeffs = rng.getrandbits(n_vars) + if planted: + affine = (coeffs & sol).bit_count() & 1 + else: + affine = rng.getrandbits(1) + eqs.append(affine | (coeffs << 1)) + return eqs, sol + + +SATISFIABLE_SHAPES = [ + pytest.param(32, 64, id="overdetermined-small"), + pytest.param(128, 200, id="overdetermined-medium"), + pytest.param(64, 64, id="square"), + pytest.param(32, 8, id="underdetermined"), +] + + +@pytest.mark.parametrize("n_vars,n_eqs", SATISFIABLE_SHAPES) +@pytest.mark.parametrize("seed", [1, 42, 2026]) +def test_solve_one_satisfies_equations(n_vars, n_eqs, seed): + eqs, _ = _random_system(n_vars, n_eqs, seed=seed) + sol = m4ri_solve(eqs, n_vars, 0) + assert sol is not None, "random consistent system must have a solution" + assert _verify(eqs, n_vars, sol) + + +@pytest.mark.parametrize("n_vars,n_eqs", SATISFIABLE_SHAPES) +@pytest.mark.parametrize("seed", [1, 42, 2026]) +def test_solve_space_origin_satisfies_equations(n_vars, n_eqs, seed): + eqs, _ = _random_system(n_vars, n_eqs, seed=seed) + space = m4ri_solve(eqs, n_vars, 1) + assert space is not None + origin = int(space.origin) + assert _verify(eqs, n_vars, origin) + + +@pytest.mark.parametrize("n_vars,n_eqs", SATISFIABLE_SHAPES) +@pytest.mark.parametrize("seed", [1, 42, 2026]) +def test_solve_one_equivalent_to_solve_space(n_vars, n_eqs, seed): + """Both paths must describe the same affine solution set. + + We require three things: + - both agree on satisfiability, + - each returned representative satisfies the system, + - `sol_one` lies in `origin ^ span(basis)` (so the fast path has not + picked a point outside the space the old path would enumerate, which + would indicate either mismatched kernels or a truncated basis). + """ + eqs, _ = _random_system(n_vars, n_eqs, seed=seed) + sol_one = m4ri_solve(eqs, n_vars, 0) + space = m4ri_solve(eqs, n_vars, 1) + assert (sol_one is None) == (space is None) + assert sol_one is not None + origin = int(space.origin) + assert _verify(eqs, n_vars, sol_one) + assert _verify(eqs, n_vars, origin) + # Every basis vector must itself lie in the kernel of the linear part. + for b in space.basis: + assert _verify(eqs, n_vars, origin ^ int(b)) + # The difference `sol_one ^ origin` must be expressible as a XOR of some + # subset of basis vectors — otherwise the two paths describe different + # affine spaces. + assert _in_span(sol_one ^ origin, space.basis) + + +@pytest.mark.parametrize("seed", [1, 42, 2026]) +def test_inconsistent_system_returns_none(seed): + # Start from a consistent random system, then flip the affine bit of one + # equation whose coefficients are all zero — giving "0 = 1". + n_vars, n_eqs = 64, 80 + eqs, _ = _random_system(n_vars, n_eqs, seed=seed) + eqs.append(1) # 0·x = 1 — never satisfiable + assert m4ri_solve(eqs, n_vars, 0) is None + assert m4ri_solve(eqs, n_vars, 1) is None + + +@pytest.mark.parametrize("seed", [1, 42, 2026]) +def test_inconsistent_via_contradicting_rows(seed): + # Pair-of-equations contradiction: same coeffs, different affine. + rng = random.Random(seed) + n_vars = 64 + coeffs = rng.getrandbits(n_vars) + eqs = [(coeffs << 1) | 0, (coeffs << 1) | 1] + assert m4ri_solve(eqs, n_vars, 0) is None + + +@pytest.mark.parametrize("seed", [1, 42, 2026]) +def test_underdetermined_space_has_expected_dimension(seed): + # 32 vars, 4 random linearly-independent-ish equations: solution space is + # almost surely 32 - rank-dimensional. We only assert dim >= 32 - n_eqs. + n_vars, n_eqs = 32, 4 + eqs, _ = _random_system(n_vars, n_eqs, seed=seed) + space = m4ri_solve(eqs, n_vars, 1) + assert space is not None + assert space.dimension >= n_vars - n_eqs + + +def test_mt19937_recovery_benchmark_case(): + """End-to-end: the canonical MT19937 state-recovery benchmark shape.""" + from gf2bv.crypto.mt import MT19937 + + rand = random.Random(3142) + state = tuple(rand.getstate()[1][:-1]) + outputs = [rand.getrandbits(32) for _ in range(624)] + + lin = LinearSystem([32] * 624) + mt = lin.gens() + rng = MT19937(mt) + zeros = [rng.getrandbits(32) ^ out for out in outputs] + [mt[0] ^ 0x80000000] + sol = lin.solve_one(zeros) + assert sol == state