error.cpp

#include "error.hpp"
#include <cstdint>
#include <stdlib.h>
#include <algorithm>
#include <iostream>


//variables set outside this file
correction_t correction_used = CORRECT_RS;
checksum_t checksum_used = CRC_8;
uint32_t correction_data = 0;
uint8_t correction_strength = 50;

const char* correction_err_msg = "";

//globals for functions for convenience
int parity_num = 0;
int sector_capacity = -1;
int checksum_size = 0;


// CRC
//we don't use any initial register value, reflection or XOR
uint8_t def_crc8 = 0x1D;
uint16_t def_crc16 = 0x8005;
uint32_t def_crc32 = 0x04C11DB7;
uint32_t crc_table[256];

void init_crc_table(uint32_t poly) {
	for (uint32_t i = 0; i < 256; i++) {
		uint32_t reg = i << 24; //shift byte to top
		for (int j = 0; j < 8; j++) {
			const bool overflow = reg >> 31;
			reg <<= 1;
			reg ^= overflow ? poly : 0;
		}
		crc_table[i] = reg;
	}
}

uint32_t get_checksum(const uint8_t *dat) {
	uint32_t crc = 0;
	const int size = sector_capacity - checksum_size;
	for (int i = 0; i < size; i++) {
		uint8_t ind = dat[i] ^ (crc >> 24);
		crc <<= 8;
		crc ^= crc_table[ind];
	}
	return crc ^ ~0;
}

void checksum(uint8_t *dat) {
	if (checksum_used == CKSUM_NONE) return;
	
	uint32_t crc = get_checksum(dat);
	//write checksum in little endian
	const int size = sector_capacity - checksum_size;
	for (int i = 0; i < checksum_size; i++)
		dat[size + i] = (crc >> (8 * (i + 4 - checksum_size))) & 0xFF;
}

//returns true if there is an error
bool verify_checksum(const uint8_t *dat) {
	if (checksum_used == CKSUM_NONE) return false;
	
	uint32_t crc = get_checksum(dat);
	//compare checksum (little endian)
	const int size = sector_capacity - checksum_size;
	for (int i = 0; i < checksum_size; i++) {
		const uint8_t crc_part = (crc >> (8 * (i + 4 - checksum_size))) & 0xFF;
		if (dat[size + i] != crc_part) return true;
	}
	return false;
}

//size given in bytes
void use_crc(uint32_t poly, int size) {
	init_crc_table(poly << (8 * (4 - size)));
	checksum_size = size;
}


// Reed Solomon
//we'll be using typical GF(2^8) and the generator polynomial method
//generator a is 2 with a^i starting at 0

//log & ilog (exp) lookup tables
uint8_t gf_log[256];
uint8_t gf_ilog[512];
uint8_t gf_prim;
uint8_t *rs_gen = nullptr; //allocated at init, freed at deinit
int rs_gen_size = 0;

void populate_rs_tables(void) {
	//use generator 2
	uint8_t g = 1;
	for (int i = 0; i < 255; i++) {
		gf_log[g] = i;
		gf_ilog[i] = g;
		
		bool overflow = (g >> 7);
		g <<= 1;
		g ^= (overflow ? gf_prim : 0);
	}
	//extend ilog table
	for (int i = 255; i < 512; i++)
		gf_ilog[i] = gf_ilog[i - 255];
}

uint8_t gf_mult(uint8_t a, uint8_t b) {
	if (a == 0 || b == 0) return 0; //no log(0)
	return gf_ilog[gf_log[a] + gf_log[b]];
}

uint8_t gf_div(uint8_t a, uint8_t b) {
	if (a == 0 || b == 0) return 0; //b == 0 is invalid
	int ind = gf_log[a];
	ind += 255 - gf_log[b]; //prevent underflow
	return gf_ilog[ind];
}

uint8_t rs_poly_eval(const uint8_t *poly, int size, uint8_t x) {
	int i = size - 1;
	uint8_t y = poly[i--];
	for (; i >= 0; i--)
		y = gf_mult(y, x) ^ poly[i];
	return y;
}

//any result that does not fit into res gets discarded
void rs_poly_mult(const uint8_t *a, const uint8_t *b, int sa, int sb, uint8_t *res, int sr) {
	for (int i = 0; i < sr; i++) res[i] = 0;
	
	for (int i = 0; i < sa; i++)
		for (int j = 0; j < sb; j++) {
			if (i + j >= sr) break;
			res[i + j] ^= gf_mult(a[i], b[j]);
		}
}

//this function implements synthetic division for monic divisors
//mod can be null and function will still compute result
//mod size must be at least sb - 1 and result size sa - sb + 1
//for sb > sa, sr >= 0, sm >= sa
void rs_poly_div(const uint8_t *a, const uint8_t *b, int sa, int sb, uint8_t *res, uint8_t *mod, int sr, int sm) {
	//actual sizes of results
	int n_sr = std::max(sa - sb + 1, 0);
	int n_sm = sa - n_sr;
	bool no_res = (res == nullptr || sr < n_sr);
	bool no_mod = (mod == nullptr || sm < n_sm);
	
	if (no_res) return;
	
	//zero out unused result part
	for (int i = n_sr; i < sr; i++) res[i] = 0;
	
	for (int i = 0; i < n_sr; i++) {
		//result has to be computed from MS to LS
		uint8_t& r = res[n_sr - 1 - i];
		r = a[sa - 1 - i]; //drop dividend element
		
		//add up diagonals aka previous results multiplied by divisor
		const int m = std::min(i, n_sm);
		for (int j = 0; j < m; j++)
			r ^= gf_mult(b[sb - 2 - j], res[n_sr - i + j]);
	}
	
	if (no_mod) return;
	
	for (int i = n_sm; i < sm; i++) mod[i] = 0;
	
	for (int i = 0; i < n_sm; i++) {
		//this time we can go LS -> MS since modulo results don't influence each other
		uint8_t& r = mod[i];
		r = a[i];
		
		//add up diagonals
		for (int j = std::max(i - n_sr + 1, 0); j <= i; j++)
			r ^= gf_mult(b[j], res[i - j]);
	}
}

void rs_generator(void) {
	rs_gen_size = parity_num + 1;
	rs_gen = (uint8_t*)calloc(rs_gen_size, sizeof(uint8_t));
	
	//we can do multiplication by (x - a^i) in place
	rs_gen[0] = 1;
	for (int i = 0; i < parity_num; i++) {
		uint8_t a = gf_ilog[i]; // 2^i
		for (int j = i; j >= 0; j--) {
			rs_gen[j+1] ^= rs_gen[j];
			rs_gen[j] = gf_mult(rs_gen[j], a);
		}
	}
}

//Berlekamp–Massey algorithm implementation
int rs_get_error_locator(uint8_t *synd, int N, uint8_t *loc) {
	const int t = parity_num / 2;
	int L = 0; //errors currently estimated
	int old_L = 0;
	int m = 1; //iterations since last L
	uint8_t old_d = 1; //discrepancy
	uint8_t old[t + 1]; //old coefficients
	
	//init coefficients
	old[0] = 1, loc[0] = 1;
	for (int i = 1; i <= t; i++) old[i] = 0, loc[i] = 0;
	
	for (int i = 0; i < N; i++) {
		//calculate discrepancy
		uint8_t d = synd[i];
		for (int k = 1; k <= L; k++) {
			if (k > i) break;
			d ^= gf_mult(loc[k], synd[i - k]);
		}
		
		if (d == 0) {
			m++;
		} else if (2 * L <= i) {
			uint8_t tmp[t + 1]; //copy of locator to replace old
			for (int i = 0; i <= L; i++) tmp[i] = loc[i];
			
			const uint8_t f = gf_div(d, old_d);
			for (int k = 0; k <= old_L; k++)
				loc[m + k] ^= gf_mult(f, old[k]);
			
			//compute new old values
			for (int i = 0; i <= L; i++) old[i] = tmp[i];
			old_L = L;
			old_d = d;
			
			L = i + 1 - L;
			if (L > t) break;
			m = 1;
		} else {
			const uint8_t f = gf_div(d, old_d);
			for (int k = 0; k <= old_L; k++)
				loc[m + k] ^= gf_mult(f, old[k]);
			m++;
		}
	}
	return L;
}

void rs_encode(uint8_t *dat) {
	const int sec_size = sector_capacity + parity_num;
	
	//prepare for division by multiplying by x^parity
	for (int i = sector_capacity - 1; i >= 0; i--)
		dat[i + parity_num] = dat[i];
	for (int i = 0; i < parity_num; i++)
		dat[i] = 0;
	
	uint8_t res[sector_capacity]; //unused for encoding
	uint8_t mod[parity_num]; //will be the result of encoding
	
	rs_poly_div(dat, rs_gen, sec_size, rs_gen_size, res, mod, sector_capacity, parity_num);
	
	//append remainder to out as parity
	for (int i = 0; i < parity_num; i++)
		dat[i] = mod[i];
}

int rs_decode(uint8_t *dat) {
	//decoding is harder than the encoding process
	//two algorithms are involved, Berlekamp-Massey and Forney
	const int sec_size = sector_capacity + parity_num;
	
	//calculate syndromes
	uint8_t synd[parity_num];
	for (int i = 0; i < parity_num; i++) {
		uint8_t a = gf_ilog[i]; // 2^i
		synd[i] = rs_poly_eval(dat, sec_size, a);
	}
	
	//check if all syndromes are 0 for shortcut
	bool no_error = true;
	for (uint8_t s : synd)
		if (s != 0) {
			no_error = false;
			break;
		}
	
	if (no_error) {
		//shift data back and remove remainder
		for (int i = 0; i < sector_capacity; i++)
			dat[i] = dat[i + parity_num];
		return 0;
	}
	
	//compute error locator
	const int t = parity_num / 2;
	uint8_t loc[t + 1];
	int L = rs_get_error_locator(synd, parity_num, loc);
	
	if (L > t) return -1; //too many errors or something else
	
	//find roots (I'm too lazy to learn Chien search)
	int err_pos[L];
	int errs = 0;
	for (int i = 0; i < 255; i++)
		if (rs_poly_eval(loc, L+1, gf_ilog[255 - i]) == 0) {
			err_pos[errs++] = i;
			if (errs >= L) break;
		}
	
	if (errs != L) return -1;
	
	//get error evaluator
	uint8_t err_eval[parity_num]; // 2t
	rs_poly_mult(synd, loc, parity_num, L+1, err_eval, parity_num);
	
	//differentiate error locator
	for (int i = 0; i < L; i++) {
		if (i % 2 != 0) {
			loc[i] = 0;
			continue;
		}
		loc[i] = loc[i + 1];
	}
	
	//finally time to get error amplitudes and correct them
	for (int i = 0; i < errs; i++) {
		const uint8_t X = gf_ilog[err_pos[i]];
		const uint8_t X_inv = gf_ilog[255 - err_pos[i]];
		const uint8_t e = gf_div(rs_poly_eval(err_eval, 2*t, X_inv), rs_poly_eval(loc, L, X_inv));
		dat[err_pos[i]] ^= gf_mult(e, X); //fix error
	}
	
	//check syndromes again
	for (int i = 0; i < parity_num; i++) {
		uint8_t a = gf_ilog[i]; // 2^i
		if (rs_poly_eval(dat, sec_size, a) != 0) return -1;
	}
	
	//shift away generator remainder
	for (int i = 0; i < sector_capacity; i++)
		dat[i] = dat[i + parity_num];
	
	return errs;
}

void init_rs(void) {
	gf_prim = correction_data & 0xFF;
	if (gf_prim == 0) gf_prim = 0x1D; //default
	populate_rs_tables();
	rs_generator();
}


int init_error_correction(int sector_size) {
	sector_capacity = -1;
	parity_num = 0;
	checksum_size = 0;
	
	switch (correction_used) {
		case CORRECT_NONE:
			//easiest case
			sector_capacity = sector_size;
			break;
		case CORRECT_RS:
			if (sector_size < 8) {
				correction_err_msg = "Sector size too small for reed solomon";
				return 8;
			}
			if (sector_size > 255) { //due to GF(2^8)
				correction_err_msg = "Sector size too big for reed solomon";
				return 255;
			}
			//evaluate graph where 100% strength -> 40% data recovery
			sector_capacity = sector_size * (100 - 4 * correction_strength / 5) / 100;
			parity_num = sector_size - sector_capacity;
			init_rs();
			break;
		default:
			correction_err_msg = "Invalid correction ID";
			return -1;
	}
	
	switch (checksum_used) {
		case CKSUM_NONE:
			break;
		case CRC_8:
			use_crc(def_crc8, 1);
			break;
		case CRC_16:
			use_crc(def_crc16, 2);
			break;
		case CRC_32:
			use_crc(def_crc32, 4);
			break;
		default:
			correction_err_msg = "Invalid checksum ID";
			return -1;
	}
	return 0;
}

void deinit_error_correction(void) {
	switch (correction_used) {
		case CORRECT_NONE:
			break;
		case CORRECT_RS:
			if (rs_gen != nullptr) free(rs_gen), rs_gen = nullptr;
			break;
	}
}

int get_sector_capacity(void) {
	return sector_capacity - checksum_size;
}

void print_ec_debug_info(void) {
	std::cout << "Sector size: " << sector_capacity << '\n';
	std::cout << "Parity size: " << parity_num << '\n';
	if (checksum_used != CKSUM_NONE)
		std::cout << "Checksum size: " << checksum_size << '\n';
	if (correction_used == CORRECT_RS) {
		std::cout << "RS generator:\n";
		for (int i = rs_gen_size - 1; i >= 0; i--)
			std::cout << (int)rs_gen[i] << ' ';
		std::cout << '\n';
	}
}

int encode(const uint8_t* in, uint8_t* out) {
	const int sec_cap = get_sector_capacity();
	for (int i = 0; i < sec_cap; i++) out[i] = in[i];
	checksum(out);
	
	switch (correction_used) {
		case CORRECT_NONE:
			break;
		case CORRECT_RS:
			rs_encode(out);
			break;
	}
	return 0;
}

int decode(uint8_t* in, uint8_t* out) {
	const int sec_cap = get_sector_capacity();
	
	int stat = 0;
	switch (correction_used) {
		case CORRECT_NONE:
			break;
		case CORRECT_RS:
			stat = rs_decode(in);
			break;
	}
	
	bool check_failed = verify_checksum(in);
	if (stat < 0 || check_failed) return -1;
	
	//move to output
	for (int i = 0; i < sec_cap; i++) out[i] = in[i];
	return stat;
}