mirror of
https://gitee.com/mirrors_PX4/PX4-Autopilot.git
synced 2026-10-25 17:33:34 +08:00
422 lines
12 KiB
Python
422 lines
12 KiB
Python
import ctypes
|
|
from enum import Enum
|
|
|
|
|
|
# Note: this implementation directly follows the heatshrink_encoder.c code
|
|
# (it's neither expected to be very efficient, nor is it pythonic)
|
|
|
|
# Enum
|
|
class HSE_state(ctypes.c_int):
|
|
HSES_NOT_FULL = 0
|
|
HSES_FILLED = 1
|
|
HSES_SEARCH = 2
|
|
HSES_YIELD_TAG_BIT = 3
|
|
HSES_YIELD_LITERAL = 4
|
|
HSES_YIELD_BR_INDEX = 5
|
|
HSES_YIELD_BR_LENGTH = 6
|
|
HSES_SAVE_BACKLOG = 7
|
|
HSES_FLUSH_BITS = 8
|
|
HSES_DONE = 9
|
|
|
|
|
|
# Constants
|
|
FLAG_IS_FINISHING = 0x01
|
|
MATCH_NOT_FOUND = 0xFFFF
|
|
|
|
HEATSHRINK_LITERAL_MARKER = 0x01
|
|
HEATSHRINK_BACKREF_MARKER = 0x00
|
|
|
|
|
|
# Structs
|
|
class output_info(ctypes.Structure):
|
|
_fields_ = [
|
|
("buf", ctypes.POINTER(ctypes.c_uint8)),
|
|
("buf_size", ctypes.c_size_t),
|
|
("output_size", ctypes.POINTER(ctypes.c_size_t))
|
|
]
|
|
|
|
|
|
# Functions
|
|
def add_tag_bit(hse, oi, tag):
|
|
push_bits(hse, 1, tag, oi)
|
|
|
|
|
|
def push_bits(hse, count, bits, oi):
|
|
assert count <= 8
|
|
current_byte = ctypes.c_uint8(hse.current_byte)
|
|
bit_index = ctypes.c_uint8(hse.bit_index)
|
|
for i in range(count - 1, -1, -1):
|
|
bit = bits & (1 << i)
|
|
if bit:
|
|
current_byte.value |= bit_index.value
|
|
bit_index.value >>= 1
|
|
if bit_index.value == 0:
|
|
bit_index.value = 0x80
|
|
oi.buf[oi.output_size[0]] = current_byte.value
|
|
oi.output_size[0] += 1
|
|
current_byte.value = 0
|
|
hse.current_byte = current_byte.value
|
|
hse.bit_index = bit_index.value
|
|
|
|
|
|
def push_literal_byte(hse, oi):
|
|
processed_offset = hse.match_scan_index - 1
|
|
input_offset = get_input_offset(hse) + processed_offset
|
|
c = hse.buffer[input_offset]
|
|
push_bits(hse, 8, c, oi)
|
|
|
|
|
|
# Define necessary structures and enums
|
|
class heatshrink_encoder(ctypes.Structure):
|
|
_fields_ = [
|
|
("input_size", ctypes.c_uint16),
|
|
("match_scan_index", ctypes.c_uint16),
|
|
("match_length", ctypes.c_uint16),
|
|
("match_pos", ctypes.c_uint16),
|
|
("outgoing_bits", ctypes.c_uint16),
|
|
("outgoing_bits_count", ctypes.c_uint8),
|
|
("flags", ctypes.c_uint8),
|
|
("state", ctypes.c_uint8),
|
|
("current_byte", ctypes.c_uint8),
|
|
("bit_index", ctypes.c_uint8),
|
|
|
|
("buffer", ctypes.POINTER(ctypes.c_uint8)),
|
|
]
|
|
|
|
def __init__(self, window_size=8, lookahead_size=4):
|
|
super().__init__()
|
|
self.window_size = window_size
|
|
self.lookahead_size = lookahead_size
|
|
self.search_index = hs_index(window_size)
|
|
self.buffer = (ctypes.c_uint8 * (2 << window_size))()
|
|
self.bit_index = 0x80
|
|
|
|
|
|
class hs_index(ctypes.Structure):
|
|
_fields_ = [("index", ctypes.POINTER(ctypes.c_int16))]
|
|
|
|
def __init__(self, window_size):
|
|
super().__init__()
|
|
self.index = (ctypes.c_int16 * (2 << window_size))()
|
|
|
|
|
|
class HSE_sink_res(Enum):
|
|
HSER_SINK_OK = 0
|
|
HSER_SINK_ERROR_NULL = -1
|
|
HSER_SINK_ERROR_MISUSE = -2
|
|
|
|
|
|
class HSE_poll_res(Enum):
|
|
HSER_POLL_EMPTY = 0
|
|
HSER_POLL_MORE = 1
|
|
HSER_POLL_ERROR_NULL = -1
|
|
HSER_POLL_ERROR_MISUSE = -2
|
|
|
|
|
|
class HSE_finish_res(Enum):
|
|
HSER_FINISH_DONE = 0
|
|
HSER_FINISH_MORE = 1
|
|
HSER_FINISH_ERROR_NULL = -1
|
|
|
|
|
|
def is_finishing(hse):
|
|
return hse.flags & FLAG_IS_FINISHING
|
|
|
|
|
|
def can_take_byte(oi):
|
|
return oi.output_size[0] < oi.buf_size
|
|
|
|
|
|
def get_input_buffer_size(hse):
|
|
return 1 << hse.window_size
|
|
|
|
|
|
def get_lookahead_size(hse):
|
|
return 1 << hse.lookahead_size
|
|
|
|
|
|
def get_input_offset(hse):
|
|
return get_input_buffer_size(hse)
|
|
|
|
|
|
def heatshrink_encoder_sink(hse, in_buf, size, input_size):
|
|
if hse is None or in_buf is None or input_size is None:
|
|
return HSE_sink_res.HSER_SINK_ERROR_NULL
|
|
|
|
if is_finishing(hse):
|
|
return HSE_sink_res.HSER_SINK_ERROR_MISUSE
|
|
|
|
if hse.state != HSE_state.HSES_NOT_FULL:
|
|
return HSE_sink_res.HSER_SINK_ERROR_MISUSE
|
|
|
|
write_offset = get_input_offset(hse) + hse.input_size
|
|
ibs = get_input_buffer_size(hse)
|
|
rem = ibs - hse.input_size
|
|
cp_sz = min(rem, size)
|
|
|
|
for i in range(cp_sz):
|
|
hse.buffer[write_offset + i] = in_buf[i]
|
|
input_size.value = cp_sz
|
|
hse.input_size += cp_sz
|
|
|
|
if cp_sz == rem:
|
|
hse.state = HSE_state.HSES_FILLED
|
|
|
|
return HSE_sink_res.HSER_SINK_OK
|
|
|
|
|
|
def do_indexing(hse):
|
|
# Build an index array I that contains flattened linked lists
|
|
# for the previous instances of every byte in the buffer.
|
|
|
|
hsi = hse.search_index
|
|
last = [0xffff] * 256
|
|
buf = hse.buffer
|
|
index = hsi.index
|
|
input_offset = get_input_offset(hse)
|
|
end = input_offset + hse.input_size
|
|
|
|
for i in range(0, end):
|
|
v = buf[i]
|
|
lv = last[v]
|
|
index[i] = lv
|
|
last[v] = i
|
|
|
|
|
|
def heatshrink_encoder_poll(hse, out_buf, out_buf_size, output_size):
|
|
if hse is None or out_buf is None or output_size is None:
|
|
return HSE_poll_res.HSER_POLL_ERROR_NULL
|
|
|
|
if out_buf_size == 0:
|
|
return HSE_poll_res.HSER_POLL_ERROR_MISUSE
|
|
|
|
output_size[0] = 0
|
|
oi = output_info()
|
|
oi.buf = out_buf
|
|
oi.buf_size = out_buf_size
|
|
oi.output_size = output_size
|
|
|
|
while True:
|
|
in_state = hse.state
|
|
if in_state == HSE_state.HSES_NOT_FULL:
|
|
return HSE_poll_res.HSER_POLL_EMPTY
|
|
elif in_state == HSE_state.HSES_DONE:
|
|
return HSE_poll_res.HSER_POLL_EMPTY
|
|
elif in_state == HSE_state.HSES_FILLED:
|
|
do_indexing(hse)
|
|
hse.state = HSE_state.HSES_SEARCH
|
|
elif in_state == HSE_state.HSES_SEARCH:
|
|
hse.state = st_step_search(hse)
|
|
elif in_state == HSE_state.HSES_YIELD_TAG_BIT:
|
|
hse.state = st_yield_tag_bit(hse, oi)
|
|
elif in_state == HSE_state.HSES_YIELD_LITERAL:
|
|
hse.state = st_yield_literal(hse, oi)
|
|
elif in_state == HSE_state.HSES_YIELD_BR_INDEX:
|
|
hse.state = st_yield_br_index(hse, oi)
|
|
elif in_state == HSE_state.HSES_YIELD_BR_LENGTH:
|
|
hse.state = st_yield_br_length(hse, oi)
|
|
elif in_state == HSE_state.HSES_SAVE_BACKLOG:
|
|
hse.state = st_save_backlog(hse)
|
|
elif in_state == HSE_state.HSES_FLUSH_BITS:
|
|
hse.state = st_flush_bit_buffer(hse, oi)
|
|
else:
|
|
return HSE_poll_res.HSER_POLL_ERROR_MISUSE
|
|
|
|
if hse.state == in_state:
|
|
if oi.output_size == oi.buf_size:
|
|
return HSE_poll_res.HSER_POLL_MORE
|
|
|
|
|
|
def heatshrink_encoder_finish(hse):
|
|
hse.flags |= FLAG_IS_FINISHING
|
|
if hse.state == HSE_state.HSES_NOT_FULL:
|
|
hse.state = HSE_state.HSES_FILLED
|
|
if hse.state == HSE_state.HSES_DONE:
|
|
return HSE_finish_res.HSER_FINISH_DONE
|
|
return HSE_finish_res.HSER_FINISH_MORE
|
|
|
|
|
|
def st_step_search(hse):
|
|
window_length = get_input_buffer_size(hse)
|
|
lookahead_sz = get_lookahead_size(hse)
|
|
msi = hse.match_scan_index
|
|
|
|
fin = is_finishing(hse)
|
|
if msi > hse.input_size - (1 if fin else lookahead_sz):
|
|
return HSE_state.HSES_FLUSH_BITS if fin else HSE_state.HSES_SAVE_BACKLOG
|
|
|
|
input_offset = get_input_offset(hse)
|
|
end = input_offset + msi
|
|
start = end - window_length
|
|
|
|
max_possible = lookahead_sz if hse.input_size - msi >= lookahead_sz else hse.input_size - msi
|
|
|
|
match_pos, match_length = find_longest_match(hse, start, end, max_possible)
|
|
|
|
if match_pos == MATCH_NOT_FOUND:
|
|
hse.match_scan_index += 1
|
|
hse.match_length = 0
|
|
return HSE_state.HSES_YIELD_TAG_BIT
|
|
else:
|
|
hse.match_pos = match_pos
|
|
hse.match_length = match_length
|
|
return HSE_state.HSES_YIELD_TAG_BIT
|
|
|
|
|
|
def find_longest_match(hse, start, end, maxlen):
|
|
buf = hse.buffer
|
|
|
|
match_maxlen = 0
|
|
match_index = MATCH_NOT_FOUND
|
|
|
|
needlepoint = end
|
|
|
|
pos = hse.search_index.index[end]
|
|
buf_needlepoint_maxlen = buf[needlepoint + match_maxlen]
|
|
while pos >= start:
|
|
|
|
if buf[pos + match_maxlen] != buf_needlepoint_maxlen:
|
|
pos = hse.search_index.index[pos]
|
|
continue
|
|
|
|
length = 1
|
|
for length in range(1, maxlen):
|
|
if buf[pos + length] != buf[needlepoint + length]:
|
|
break
|
|
if length > match_maxlen:
|
|
match_maxlen = length
|
|
match_index = pos
|
|
buf_needlepoint_maxlen = buf[needlepoint + match_maxlen]
|
|
if length == maxlen:
|
|
break # won't find better
|
|
|
|
pos = hse.search_index.index[pos]
|
|
|
|
break_even_point = 1 + hse.window_size + hse.lookahead_size
|
|
|
|
if match_maxlen > (break_even_point // 8):
|
|
return end - match_index, match_maxlen
|
|
return MATCH_NOT_FOUND, 0
|
|
|
|
|
|
def push_outgoing_bits(hse, oi):
|
|
if hse.outgoing_bits_count > 8:
|
|
count = 8
|
|
bits = hse.outgoing_bits >> (hse.outgoing_bits_count - 8)
|
|
else:
|
|
count = hse.outgoing_bits_count
|
|
bits = hse.outgoing_bits
|
|
|
|
if count > 0:
|
|
push_bits(hse, count, bits, oi)
|
|
hse.outgoing_bits_count -= count
|
|
return count
|
|
|
|
|
|
def st_yield_tag_bit(hse, oi):
|
|
if can_take_byte(oi):
|
|
if hse.match_length == 0:
|
|
add_tag_bit(hse, oi, HEATSHRINK_LITERAL_MARKER)
|
|
return HSE_state.HSES_YIELD_LITERAL
|
|
else:
|
|
add_tag_bit(hse, oi, HEATSHRINK_BACKREF_MARKER)
|
|
hse.outgoing_bits = hse.match_pos - 1
|
|
hse.outgoing_bits_count = hse.window_size
|
|
return HSE_state.HSES_YIELD_BR_INDEX
|
|
else:
|
|
return HSE_state.HSES_YIELD_TAG_BIT
|
|
|
|
|
|
def st_yield_literal(hse, oi):
|
|
if can_take_byte(oi):
|
|
push_literal_byte(hse, oi)
|
|
return HSE_state.HSES_SEARCH
|
|
else:
|
|
return HSE_state.HSES_YIELD_LITERAL
|
|
|
|
|
|
def st_yield_br_index(hse, oi):
|
|
if can_take_byte(oi):
|
|
if push_outgoing_bits(hse, oi) > 0:
|
|
return HSE_state.HSES_YIELD_BR_INDEX
|
|
else:
|
|
hse.outgoing_bits = hse.match_length - 1
|
|
hse.outgoing_bits_count = hse.lookahead_size
|
|
return HSE_state.HSES_YIELD_BR_LENGTH
|
|
else:
|
|
return HSE_state.HSES_YIELD_BR_INDEX
|
|
|
|
|
|
def st_yield_br_length(hse, oi):
|
|
if can_take_byte(oi):
|
|
if push_outgoing_bits(hse, oi) > 0:
|
|
return HSE_state.HSES_YIELD_BR_LENGTH
|
|
else:
|
|
hse.match_scan_index += hse.match_length
|
|
hse.match_length = 0
|
|
return HSE_state.HSES_SEARCH
|
|
else:
|
|
return HSE_state.HSES_YIELD_BR_LENGTH
|
|
|
|
|
|
def st_save_backlog(hse):
|
|
save_backlog(hse)
|
|
return HSE_state.HSES_NOT_FULL
|
|
|
|
|
|
def st_flush_bit_buffer(hse, oi):
|
|
if hse.bit_index == 0x80:
|
|
return HSE_state.HSES_DONE
|
|
elif can_take_byte(oi):
|
|
oi.buf[oi.output_size[0]] = hse.current_byte
|
|
oi.output_size[0] += 1
|
|
return HSE_state.HSES_DONE
|
|
else:
|
|
return HSE_state.HSES_FLUSH_BITS
|
|
|
|
|
|
def save_backlog(hse):
|
|
input_buf_sz = get_input_buffer_size(hse)
|
|
|
|
msi = hse.match_scan_index
|
|
|
|
rem = input_buf_sz - msi # unprocessed bytes
|
|
shift_sz = input_buf_sz + rem
|
|
|
|
for i in range(shift_sz):
|
|
hse.buffer[i] = hse.buffer[input_buf_sz - rem + i]
|
|
|
|
hse.match_scan_index = 0
|
|
hse.input_size -= input_buf_sz - rem
|
|
|
|
|
|
def encode(data, window_size, lookahead_size):
|
|
hse = heatshrink_encoder(window_size, lookahead_size)
|
|
input_buf = (ctypes.c_uint8 * len(data))()
|
|
for i, d in enumerate(data):
|
|
input_buf[i] = d
|
|
in_size = len(input_buf)
|
|
|
|
out_buf_size = 4 * in_size # set output buffer size a bit larger
|
|
out_buf = (ctypes.c_uint8 * out_buf_size)()
|
|
|
|
sunk = 0
|
|
ret = []
|
|
while sunk < in_size:
|
|
input_size = ctypes.c_size_t(in_size)
|
|
heatshrink_encoder_sink(hse, input_buf, in_size - sunk, input_size)
|
|
input_buf = input_buf[input_size.value:]
|
|
sunk += input_size.value
|
|
if sunk == in_size:
|
|
heatshrink_encoder_finish(hse)
|
|
|
|
poll_res = HSE_poll_res.HSER_POLL_MORE
|
|
while poll_res == HSE_poll_res.HSER_POLL_MORE:
|
|
output_size = (ctypes.c_size_t * 1)()
|
|
poll_res = heatshrink_encoder_poll(hse, out_buf, out_buf_size, output_size)
|
|
ret += list(out_buf)[0:output_size[0]]
|
|
|
|
if sunk == in_size:
|
|
heatshrink_encoder_finish(hse)
|
|
return ret
|