research/erasure_code/share.h

550 lines
17 KiB
C++

#include <array>
#include <iostream>
#include <exception>
#include <cassert>
#include <cstdint>
#include <vector>
#include "utils.h"
class ZeroDivisionError : std::domain_error {
public:
ZeroDivisionError() : domain_error("division by zero") { }
};
// GF(2^8) in the form (Z/2Z)[x]/(x^8+x^4+x^3+x+1)
// (the AES polynomial)
class Galois {
// the coefficients of the polynomial, where the ith bit of `val` is the x^i
// coefficient
std::uint8_t v;
// precomputed data: log and exp tables
static const std::array<Galois, 255> exptable;
static const std::array<std::uint8_t, 256> logtable;
public:
explicit constexpr Galois(unsigned char val) : v(val) { }
Galois operator+(Galois b) const {
return Galois(v ^ b.v);
}
Galois operator-(Galois b) const {
return Galois(v ^ b.v);
}
Galois operator*(Galois b) const {
return v == 0 || b.v == 0
? Galois(0)
: exptable[(unsigned(logtable[v]) + logtable[b.v]) % 255];
}
Galois operator/(Galois b) const {
if (b.v == 0) {
throw ZeroDivisionError();
}
return v == 0 || b.v == 0
? Galois(0)
: exptable[(unsigned(logtable[v]) + 255u - logtable[b.v]) % 255];
}
Galois operator-() const {
return *this;
}
Galois& operator+=(Galois b) {
return *this = *this + b;
}
Galois& operator-=(Galois b) {
return *this = *this - b;
}
Galois& operator*=(Galois b) {
return *this = *this * b;
}
Galois& operator/=(Galois b) {
return *this = *this / b;
}
bool operator==(Galois b) {
return v == b.v;
}
// back door
std::uint8_t val() const {
return v;
}
};
// Z/pZ, for an odd prime p
template<unsigned p>
class Modulo {
// check that p is prime by trial division
static constexpr bool is_prime(unsigned x, unsigned divisor = 2) {
return divisor * divisor > x
? true
: x % divisor != 0 && is_prime(x, divisor + 1);
}
static_assert(p > 2 && is_prime(p, 2), "p must be an odd prime!");
unsigned v;
public:
explicit Modulo(unsigned val) : v(val) {
assert(v >= 0 && v < p);
}
Modulo inv() const {
if (v == 0) {
throw ZeroDivisionError();
}
unsigned r = 1, base = v, exp = p-2;
while (exp > 0) {
if (exp & 1) r = (r * base) % p;
base = (base * base) % p;
exp >>= 1;
}
return Modulo(r);
}
Modulo operator+(Modulo b) const {
return Modulo((v + b.v) % p);
}
Modulo operator-(Modulo b) const {
return Modulo((v + p - b.v) % p);
}
Modulo operator*(Modulo b) const {
return Modulo((v * b.v) % p);
}
Modulo operator/(Modulo b) const {
return *this * b.inv();
}
Modulo& operator+=(Modulo b) {
return *this = *this + b;
}
Modulo& operator-=(Modulo b) {
return *this = *this - b;
}
Modulo& operator*=(Modulo b) {
return *this = *this * b;
}
Modulo& operator/=(Modulo b) {
return *this = *this / b;
}
bool operator==(Modulo b) {
return v == b.v;
}
// back door
unsigned val() const {
return v;
}
};
// Evaluates a polynomial p in little-endian form (e.g. x^2 + 3x + 2 is
// represented as {2, 3, 1}) at coordinate x,
// e.g. eval_poly_at((int[]){2, 3, 1}, 5) = 42.
//
// T should be a type supporting ring arithmetic and T(0) and T(1) should be the
// appropriate identities.
//
// Range should be a type that can be iterated to get const T& elements.
template<typename T, typename Range>
T eval_poly_at(const Range& p, T x) {
T r(0), xi(1);
for (const T& c_i : p) {
r += c_i * xi;
xi *= x;
}
return r;
}
// Given p+1 y values and x values with no errors, recovers the original
// degree-p polynomial. For example,
// lagrange_interp<double>((double[]){51.0, 59.0, 66.0},
// (double[]){1.0, 3.0, 4.0})
// = {50.0, 0.0, 1.0}.
//
// T should be a field and Range should be a sized range type with values of
// type T. T(0) and T(1) should be the appropriate field identities.
template<typename T, typename Range>
std::vector<T> lagrange_interp(const Range& pieces, const Range& xs) {
// `size` is the number of datapoints; the degree of the result polynomial
// is then `size-1`
const unsigned size = pieces.size();
assert(size == xs.size());
std::vector<T> root{T(1)}; // initially just the polynomial "1"
// build up the numerator polynomial, `root`, by taking the product of (x-v)
// (implemented as convolving repeatedly with [-v, 1])
for (const T& v : xs) {
// iterate backward since new root[i] depends on old root[i-1]
for (unsigned i = root.size(); i--; ) {
root[i] *= -v;
if (i > 0) root[i] += root[i-1];
}
// polynomial is always monic so save an extra multiply by doing this
// after
root.emplace_back(1);
}
// should have degree `size`
assert(root.size() == size + 1);
// generate per-value numerator polynomials by dividing the master
// polynomial back by each x coordinate
std::vector<std::vector<T> > nums;
nums.reserve(size);
for (const T& v : xs) {
// divide `root` by (x-v) to get a degree size-1 polynomial
// (i.e. with `size` coefficients)
std::vector<T> num(size, T(0));
// compute the x^0, x^1, ..., x^(p-2) coefficients by long division
T last = num.back() = T(1); // still always a monic polynomial
for (int i = int(size)-2; i >= 0; --i) {
num[i] = last = root[i+1] + last * v;
}
nums.emplace_back(std::move(num));
}
assert(nums.size() == size);
// generate denominators by evaluating numerator polys at their x
std::vector<T> denoms;
denoms.reserve(size);
{
unsigned i = 0;
for (const T& v : xs) {
denoms.push_back(eval_poly_at(nums[i], v));
++i;
}
}
assert(denoms.size() == size);
// generate output polynomial by taking the sum over i of
// (nums[i] * pieces[i] / denoms[i])
std::vector<T> sum(size, T(0));
{
unsigned i = 0;
for (const T& y : pieces) {
T factor = y / denoms[i];
// add nums[i] * factor to sum, as a vector
for (unsigned j = 0; j < size; ++j) {
sum[j] += nums[i][j] * factor;
}
++i;
}
}
return sum;
}
// Given two linear equations, eliminates the first variable and returns
// the resulting equation.
//
// An equation of the form a_1 x_1 + ... + a_n x_n + b = 0
// is represented as the array [a_1, ..., a_n, b].
//
// T should be a ring and Range should be an indexable, sized range of T.
template<typename T, typename Range>
std::vector<T> elim(const Range& a, const Range& b) {
assert(a.size() == b.size());
std::vector<T> result;
const unsigned size = a.size();
for (unsigned i = 1; i < size; ++i) {
result.push_back(a[i] * b[0] - b[i] * a[0]);
}
return result;
}
// Given one homogeneous linear equation and the values of all but the first
// variable, solve for the value of the first variable.
//
// For an equation of the form
// a_1 x_1 + ... + a_n x_n = 0
// pass two arrays, [a_1, ..., a_n] and [x_2, ..., x_n].
//
// T should be a field; and R1 and R2 should be indexable, sized ranges of T.
template<typename T, typename R1, typename R2>
T evaluate(const R1& coeffs, const R2& vals) {
assert(coeffs.size() == vals.size() + 1);
T total(0);
const unsigned size = vals.size();
for (unsigned i = 0; i < size; ++i) {
total -= coeffs[i+1] * vals[i];
}
return total / coeffs[0];
}
// Given an n*n system of inhomogeneous linear equations, solve for the value of
// every variable.
//
// For equations of the form
// a_1,1 x_1 + ... + a_1,n x_n + b_1 = 0
// a_2,1 x_1 + ... + a_2,n x_n + b_2 = 0
// ...
// a_n,1 x_1 + ... + a_n,n x_n + b_n = 0
// pass a two-dimensional array
// [[a_1,1, ..., a_1,n, b_1], ..., [a_n,1, ..., a_n,n, b_n]].
//
// Returns the values of [x_1, ..., x_n].
//
// T should be a field.
template<typename T>
std::vector<T> sys_solve(std::vector<std::vector<T>> eqs) {
assert(eqs.size() > 0);
std::vector<std::vector<T>> back_eqs{eqs[0]};
while (eqs.size() > 1) {
std::vector<std::vector<T>> neweqs;
neweqs.reserve(eqs.size()-1);
for (unsigned i = 0; i < eqs.size()-1; ++i) {
neweqs.push_back(elim<T>(eqs[i], eqs[i+1]));
}
eqs = std::move(neweqs);
// find a row with a nonzero first entry
unsigned i = 0;
while (i + 1 < eqs.size() && eqs[i][0] == T(0)) {
++i;
}
back_eqs.push_back(eqs[i]);
}
std::vector<T> kvals(back_eqs.size()+1, T(0));
kvals.back() = T(1);
// back-substitute in reverse order
// (smallest to largest equation)
for (unsigned i = back_eqs.size(); i--; ) {
kvals[i] = evaluate<T>(back_eqs[i],
// use the already-computed values + the 1 at the end
make_iter_pair(kvals.begin()+i+1, kvals.end()));
}
kvals.pop_back();
return kvals;
}
// Divide two polynomials with nonzero leading terms.
// T should be a field.
template<typename T>
std::vector<T> polydiv(std::vector<T> Q, const std::vector<T>& E) {
if (Q.size() < E.size()) return {};
std::vector<T> div(Q.size() - E.size() + 1, T(0));
unsigned i = div.size();
while (i--) {
T factor = Q.back() / E.back();
div[i] = factor;
// subtract factor * E * x^i from Q
Q.pop_back(); // the highest term should cancel
for (unsigned j = 0; j < E.size() - 1; ++j) {
Q[i+j] -= factor * E[j];
}
assert(Q.size() == i + E.size() - 1);
}
return div;
}
// Given a set of y coordinates and x coordinates, and the degree of the
// original polynomial, determines the original polynomial even if some of the y
// coordinates are wrong. If m is the minimal number of pieces (ie. degree +
// 1), t is the total number of pieces provided, then the algo can handle up to
// (t-m)/2 errors.
//
// T should be a field. In particular, division by zero over T should throw
// ZeroDivisionError.
template<typename T>
std::vector<T> berlekamp_welch_attempt(const std::vector<T>& pieces,
const std::vector<T>& xs, unsigned master_degree) {
const unsigned error_locator_degree = (pieces.size() - master_degree - 1) / 2;
// Set up the equations for y[i]E(x[i]) = Q(x[i])
// degree(E) = error_locator_degree
// degree(Q) = master_degree + error_locator_degree - 1
std::vector<std::vector<T>> eqs(2*error_locator_degree + master_degree + 1);
for (unsigned i = 0; i < eqs.size(); ++i) {
std::vector<T>& eq = eqs[i];
const T& x = xs[i];
const T& piece = pieces[i];
T neg_x_j = T(0) - T(1);
for (unsigned j = 0; j < error_locator_degree + master_degree + 1; ++j) {
eq.push_back(neg_x_j);
neg_x_j *= x;
}
T x_j = T(1);
for (unsigned j = 0; j < error_locator_degree + 1; ++j) {
eq.push_back(x_j * piece);
x_j *= x;
}
}
// Solve the equations
// Assume the top error polynomial term to be one
int errors = error_locator_degree;
unsigned ones = 1;
std::vector<T> polys;
while (errors >= 0) {
try {
polys = sys_solve(eqs);
} catch (const ZeroDivisionError&) {
eqs.pop_back();
for (auto& eq : eqs) {
eq[eq.size()-2] += eq.back();
eq.pop_back();
}
--errors;
++ones;
continue;
}
for (unsigned i = 0; i < ones; ++i) polys.emplace_back(1);
break;
}
if (errors < 0) {
throw std::logic_error("Not enough data!");
}
// divide the polynomials...
const unsigned split = error_locator_degree + master_degree + 1;
std::vector<T> div = polydiv(std::vector<T>(polys.begin(), polys.begin() + split),
std::vector<T>(polys.begin() + split, polys.end()));
unsigned corrects = 0;
for (unsigned i = 0; i < xs.size(); ++i) {
if (eval_poly_at<T>(div, xs[i]) == pieces[i]) {
++corrects;
}
}
if (corrects < master_degree + errors) {
throw std::logic_error("Answer doesn't match (too many errors)!");
}
return div;
}
// Extends a list of integers in [0 ... 255] (if using Galois arithmetic) by
// adding n redundant error-correction values
template<typename T, typename F=Galois>
std::vector<T> extend(std::vector<T> data, unsigned n) {
const unsigned size = data.size();
std::vector<F> data_f;
data_f.reserve(size);
for (T d : data) data_f.emplace_back(d);
std::vector<F> xs;
for (unsigned i = 0; i < size; ++i) xs.emplace_back(i);
std::vector<F> poly = berlekamp_welch_attempt(data_f, xs, size-1);
data.reserve(size+n);
for (unsigned i = 0; i < n; ++i) {
data.push_back(eval_poly_at(poly, F(T(size + i))).val());
}
return data;
}
// Repairs a list of integers in [0 ... 255]. Some integers can be erroneous,
// and you can put -1 in place of an integer if you know that a certain
// value is defective or missing. Uses the Berlekamp-Welch algorithm to
// do error-correction
template<typename T, typename F=Galois>
std::vector<T> repair(const std::vector<T>& data, unsigned datasize) {
std::vector<F> vs, xs;
for (unsigned i = 0; i < data.size(); ++i) {
if (data[i] >= 0) {
vs.emplace_back(data[i]);
xs.emplace_back(T(i));
}
}
std::vector<F> poly = berlekamp_welch_attempt(vs, xs, datasize - 1);
std::vector<T> result;
for (unsigned i = 0; i < data.size(); ++i) {
result.push_back(eval_poly_at(poly, F(T(i))).val());
}
return result;
}
template<typename T>
std::vector<std::vector<T>> transpose(const std::vector<std::vector<T>>& d) {
assert(d.size() > 0);
unsigned width = d[0].size();
std::vector<std::vector<T>> result(width);
for (unsigned i = 0; i < width; ++i) {
for (unsigned j = 0; j < d.size(); ++j) {
result[i].push_back(d[j][i]);
}
}
return result;
}
template<typename T>
std::vector<T> extract_column(const std::vector<std::vector<T>>& d, unsigned i) {
std::vector<T> result;
for (unsigned j = 0; j < d.size(); ++j) {
result.push_back(d[j][i]);
}
return result;
}
// Extends a list of bytearrays
// eg. extend_chunks([map(ord, 'hello'), map(ord, 'world')], 2)
// n is the number of redundant error-correction chunks to add
template<typename T, typename F=Galois>
std::vector<std::vector<T>> extend_chunks(
const std::vector<std::vector<T>>& data,
unsigned n) {
std::vector<std::vector<T>> o;
const unsigned height = data.size();
assert(height > 0);
const unsigned width = data[0].size();
for (unsigned i = 0; i < width; ++i) {
o.push_back(extend<T, F>(extract_column(data, i), n));
}
return transpose(o);
}
// Repairs a list of bytearrays. Use an empty array in place of a missing array.
// Individual arrays can contain some missing or erroneous data.
template<typename T, typename F=Galois>
std::vector<std::vector<T>> repair_chunks(
std::vector<std::vector<T>> data,
unsigned datasize) {
unsigned width = 0;
for (const std::vector<T>& row : data) {
if (row.size() > 0) {
width = row.size();
break;
}
}
assert(width > 0);
for (std::vector<T>& row : data) {
if (row.size() == 0) {
row.assign(width, -1);
} else {
assert(row.size() == width);
}
}
std::vector<std::vector<T>> o;
for (unsigned i = 0; i < width; ++i) {
o.push_back(repair<T, F>(extract_column(data, i), datasize));
}
return transpose(o);
}
// Extends either a bytearray or a list of bytearrays or a list of lists...
// Used in the cubify method to expand a cube in all dimensions
template<typename T, typename F=Galois>
struct deep_extend_chunks_helper {
static std::vector<T> go(const std::vector<T>& data, unsigned n) {
return extend<T, Galois>(data, n);
}
};
template<typename T, typename F>
struct deep_extend_chunks_helper<std::vector<T>, F> {
static std::vector<std::vector<T>> go(const std::vector<std::vector<T>>& data, unsigned n) {
std::vector<std::vector<T>> o;
const unsigned height = data.size();
assert(height > 0);
const unsigned width = data[0].size();
for (unsigned i = 0; i < width; ++i) {
o.push_back(deep_extend_chunks_helper<T, F>::go(extract_column(data, i), n));
}
return transpose(o);
}
};
template<typename T, typename F=Galois>
std::vector<T> deep_extend_chunks(const std::vector<T>& data, unsigned n) {
return deep_extend_chunks_helper<T, F>::go(data, n);
}