144 lines
3.6 KiB
Python
144 lines
3.6 KiB
Python
from enum import Enum, auto
|
|
|
|
from dataclasses import dataclass
|
|
|
|
def __bit_array(c):
|
|
assert c <= 255 and c >= 0, c
|
|
for i in range(7, -1, -1):
|
|
yield (c >> i) & 1
|
|
|
|
def _bit_array(b):
|
|
for c in b:
|
|
yield from __bit_array(c)
|
|
|
|
def bit_array(b):
|
|
return list(_bit_array(b))
|
|
|
|
def data_bit_pred(pn, n):
|
|
return (n & (1 << (pn - 1))) != 0
|
|
|
|
class ParityResult(Enum):
|
|
zero_bit = auto()
|
|
one_bit = auto()
|
|
two_bit = auto()
|
|
|
|
def __repr__(self):
|
|
cls_name = self.__class__.__name__
|
|
return f'{cls_name}.{self.name}'
|
|
|
|
@dataclass
|
|
class HammingCode:
|
|
bits: int
|
|
|
|
@property
|
|
def block_length(self):
|
|
return 2 ** self.bits - 1
|
|
|
|
@property
|
|
def message_length(self):
|
|
return 2 ** self.bits - self.bits - 1
|
|
|
|
def all_data_bits(self):
|
|
power = 0
|
|
for bit_ix in range(1, self.block_length + 1):
|
|
if bit_ix == 2 ** power:
|
|
power += 1
|
|
else:
|
|
yield bit_ix
|
|
|
|
def interleave_parity(self, ba):
|
|
assert len(ba) == self.message_length
|
|
power = 0
|
|
index = 0
|
|
yield None # extended parity
|
|
for bit_ix in range(1, self.block_length + 1):
|
|
if bit_ix == 2 ** power:
|
|
yield None
|
|
power += 1
|
|
else:
|
|
yield ba[index]
|
|
index += 1
|
|
|
|
def parity_data_bits(self, power):
|
|
mask = 2 ** power
|
|
for bit_ix in range(1, self.block_length + 1):
|
|
if bit_ix & mask:
|
|
yield bit_ix
|
|
|
|
def parity(self, pba, power):
|
|
p = sum(
|
|
pba[bit_ix]
|
|
for bit_ix in self.parity_data_bits(power)
|
|
if bit_ix != 2 ** power
|
|
)
|
|
return p % 2
|
|
|
|
def extended_parity(self, pba):
|
|
p = sum(
|
|
pba[bit_ix]
|
|
for bit_ix in range(1, self.block_length + 1)
|
|
)
|
|
return p % 2
|
|
|
|
def encode_bitarray(self, ba):
|
|
pba = list(self.interleave_parity(ba))
|
|
assert len(pba) == self.block_length + 1
|
|
for power in range(self.bits):
|
|
parity_ix = 2 ** power
|
|
assert pba[parity_ix] is None
|
|
pba[parity_ix] = self.parity(pba, power)
|
|
pba[0] = self.extended_parity(pba)
|
|
return pba
|
|
|
|
def encode_block(self, b):
|
|
assert len(b) * 8 == self.message_length, (len(b), self.message_length)
|
|
ba = bit_array(b)
|
|
return self.encode_bitarray(ba)
|
|
|
|
def _check_parity(self, pba):
|
|
for power in range(self.bits):
|
|
parity_ix = 2 ** power
|
|
yield self.parity(pba, power) == pba[parity_ix]
|
|
|
|
def check_parity(self, pba):
|
|
error_bits = set(self.all_data_bits())
|
|
errors = 0
|
|
for power, match in enumerate(self._check_parity(pba)):
|
|
data_bits = set(self.parity_data_bits(power))
|
|
if match:
|
|
error_bits -= data_bits
|
|
else:
|
|
errors += 1
|
|
error_bits &= data_bits
|
|
if errors > 0:
|
|
if self.extended_parity(pba) == pba[0]:
|
|
return ParityResult.two_bit, error_bits
|
|
else:
|
|
assert len(error_bits) == 1, error_bits
|
|
return ParityResult.one_bit, error_bits
|
|
else:
|
|
assert len(error_bits) == 0, error_bits
|
|
return ParityResult.zero_bit, error_bits
|
|
|
|
hc = HammingCode(bits=4)
|
|
pba = hc.encode_bitarray([
|
|
1,
|
|
1, 0, 0,
|
|
1, 0, 1,
|
|
1, 0, 1, 1
|
|
])
|
|
|
|
for y in range(4):
|
|
for x in range(4):
|
|
bit = pba[y * 4 + x]
|
|
print(bit, end=' ')
|
|
print()
|
|
|
|
pba = [
|
|
1, 1, 0, 1,
|
|
0, 1, 0, 0,
|
|
1, 1, 0, 1,
|
|
1, 0, 1, 1,
|
|
]
|
|
print(hc.check_parity(pba))
|