code/mersenne_packing.hh at master · aicodix/code · GitHub

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
/*
Mersenne 2^31-1 Prime Field Packing

Copyright 2026 Ahmet Inan <inan@aicodix.de>
*/

#pragma once

#include "prime_field.hh"

namespace CODE {

struct MersennePacking
{
	typedef PrimeField<uint32_t, 0x7FFFFFFF> M31;
	static void pack(M31 *dst, const uint8_t *src, int count, int64_t bytes)
	{
		uint64_t acc = 0;
		int64_t pos = 0;
		for (int i = 0, k = 0; i < count; i++) {
			while (k < 31 && pos < bytes) {
				acc |= uint64_t(src[pos++]) << k;
				k += 8;
			}
			dst[i] = M31(acc & 0x7FFFFFFF);
			acc >>= 31;
			k -= 31;
		}
	}
	static void unpack(uint8_t *dst, const M31 *src, int count, int64_t bytes)
	{
		uint64_t acc = 0;
		int64_t pos = 0;
		for (int i = 0, k = 0; i < count; i++) {
			acc |= uint64_t(src[i].v) << k;
			k += 31;
			while ((k >= 8 || i == count - 1) && pos < bytes) {
				dst[pos++] = acc & 255;
				acc >>= 8;
				k -= 8;
			}
		}
	}
	static void decode(uint8_t *dst, const M31 *src, int64_t bytes, M31 sub)
	{
		int count = (bytes * 8 + 30) / 31;
		uint64_t acc = 0;
		int64_t pos = 0;
		for (int i = 0, k = 0; i < count; i++) {
			uint64_t val = sub == src[i] ? 0x7FFFFFFF : src[i]();
			acc |= val << k;
			k += 31;
			while (k >= 8 && pos < bytes) {
				dst[pos++] = acc & 255;
				acc >>= 8;
				k -= 8;
			}
		}
	}
};

template <int64_t MAX_BYTES>
struct MersenneRemapping
{
	typedef PrimeField<uint32_t, 0x7FFFFFFF> M31;
	typedef unsigned used_word;
	static constexpr int max_count = (MAX_BYTES * 8 + 30) / 31;
	static_assert(max_count < 0x7FFFFFFF, "Block length must be smaller than number of field values");
	static constexpr int used_width = 8 * sizeof(used_word);
	static constexpr int used_length = max_count / used_width + 1;
	used_word used_values[used_length];
	MersenneRemapping(){}
	M31 find_unused(const M31 *dst, int count)
	{
		int limit = count / used_width + 1;
		for (int i = 0; i < limit; ++i)
			used_values[i] = 0;
		for (int i = 0; i < count; ++i)
			if (int(dst[i].v)/used_width < limit)
				used_values[dst[i].v/used_width] |= 1 << dst[i].v%used_width;
		int s = 0;
		while (s/used_width < limit && used_values[s/used_width] & 1 << s%used_width)
			++s;
		return M31(s);
	}
	M31 encode(M31 *dst, const uint8_t *src, int64_t bytes)
	{
		assert(bytes <= MAX_BYTES);
		int count = (bytes * 8 + 30) / 31;
		MersennePacking::pack(dst, src, count, bytes);
		M31 sub = find_unused(dst, count);
		for (int i = 0; i < count; ++i)
			if (dst[i].v == 0x7FFFFFFF)
				dst[i] = sub;
		return sub;
	}
};

}