///|
// OpenEXR PIZ uses Haar lifting on each 16-bit lane of a channel.
// Short ranges permit signed differences; wider ranges lift modulo 2^16.
fn piz_pair(
a : Int32,
b : Int32,
narrow : Bool,
inverse : Bool,
) -> (Int32, Int32) {
if narrow {
let a = a << 16 >> 16
let b = b << 16 >> 16
if inverse {
let left = a + ((b + 1) >> 1)
(left & 65535, (left - b) & 65535)
} else {
(((a + b) >> 1) & 65535, (a - b) & 65535)
}
} else if inverse {
let right = (a - (b >> 1)) & 65535
((b + right - 32768) & 65535, right)
} else {
let left = a ^ 32768
let delta = left - b
let mean = (left + b) >> 1
(if delta < 0 { mean ^ 32768 } else { mean }, delta & 65535)
}
}
///|
fn piz_pair_at(
words : Array[Int32],
a : Int32,
b : Int32,
narrow : Bool,
inverse : Bool,
) -> Unit {
let (left, right) = piz_pair(words[a], words[b], narrow, inverse)
words[a] = left
words[b] = right
}
///|
fn piz_wavelet_step(
words : Array[Int32],
start : Int32,
width : Int32,
height : Int32,
lanes : Int32,
step : Int32,
narrow : Bool,
inverse : Bool,
) -> Unit {
let span = step * 2
let dx = step * lanes
let dy = dx * width
let mut y = 0
while y <= height - span {
let row = start + y * width * lanes
let mut x = 0
while x <= width - span {
let a = row + x * lanes
let b = a + dx
let c = a + dy
let d = c + dx
if inverse {
piz_pair_at(words, a, c, narrow, true)
piz_pair_at(words, b, d, narrow, true)
}
piz_pair_at(words, a, b, narrow, inverse)
piz_pair_at(words, c, d, narrow, inverse)
if !inverse {
piz_pair_at(words, a, c, narrow, false)
piz_pair_at(words, b, d, narrow, false)
}
x += span
}
if (width & step) != 0 {
let a = row + x * lanes
piz_pair_at(words, a, a + dy, narrow, inverse)
}
y += span
}
if (height & step) != 0 {
let row = start + y * width * lanes
for x = 0; x <= width - span; x = x + span {
let a = row + x * lanes
piz_pair_at(words, a, a + dx, narrow, inverse)
}
}
}
///|
fn piz_wavelet(
words : Array[Int32],
channels : Array[Channel],
window : Window,
max_value : Int32,
inverse : Bool,
) -> Unit {
let width = window.width()
let height = window.height()
let n = imin(width, height)
let mut start = 0
for ch in channels {
let lanes = ch.kind.bytes() / 2
for lane = 0; lane < lanes; lane = lane + 1 {
let mut step = 1
if inverse {
while step <= n / 2 {
step *= 2
}
step /= 2
}
while step >= 1 && step <= n / 2 {
piz_wavelet_step(
words,
start + lane,
width,
height,
lanes,
step,
max_value < 16384,
inverse,
)
step = if inverse { step / 2 } else { step * 2 }
}
}
start += width * height * lanes
}
}
///|
// The EXR chunk interleaves channel rows, while the transform needs channel planes.
fn piz_positions(channels : Array[Channel], w : Window) -> Array[Int32] {
let mut start = 0
channels.map(ch => {
let pos = start
start += w.width() * w.height() * (ch.kind.bytes() / 2)
pos
})
}
///|
fn piz_pack(
raw : Bytes,
channels : Array[Channel],
w : Window,
) -> Bytes raise ExrError {
if raw.length() <= 29 {
return raw
}
let positions = piz_positions(channels, w)
let words = Array::make(raw.length() / 2, 0)
let bitmap = Array::make(8192, 0)
let mut pos = 0
let mut first = 8191
let mut last = 0
for _row = 0; _row < w.height(); _row = _row + 1 {
for i = 0; i < channels.length(); i = i + 1 {
let row_words = w.width() * (channels[i].kind.bytes() / 2)
for j = 0; j < row_words; j = j + 1 {
let word = raw[pos].to_int() | (raw[pos + 1].to_int() << 8)
pos += 2
words[positions[i] + j] = word
if word != 0 {
let entry = word >> 3
bitmap[entry] = bitmap[entry] | (1 << (word & 7))
first = imin(first, entry)
last = imax(last, entry)
}
}
positions[i] += row_words
}
}
let span = imax(0, last - first + 1)
let budget = raw.length() - 9 - span
if budget <= 20 {
return raw
}
let lut = Array::make((last + 1) * 8, 0)
let mut rank = 0
for value = imax(1, first * 8); value < lut.length(); value = value + 1 {
if (bitmap[value >> 3] & (1 << (value & 7))) != 0 {
rank += 1
lut[value] = rank
}
}
for i = 0; i < words.length(); i = i + 1 {
words[i] = lut[words[i]]
}
piz_wavelet(words, channels, w, rank, false)
match huffman_encode(words, budget) {
None => raw
Some(encoded) => {
let out = sink()
out.uint(first.to_uint64(), 2)
out.uint(last.to_uint64(), 2)
for i = first; i <= last; i = i + 1 {
out.bytes.push(bitmap[i].to_byte())
}
out.i32(encoded.length())
out.raw(encoded)
out.finish()
}
}
}
///|
fn piz_unpack(
bytes : Bytes,
channels : Array[Channel],
w : Window,
expected : Int32,
) -> Bytes raise ExrError {
let r = reader(bytes)
let first = r.uint(2).to_int()
let last = r.uint(2).to_int()
if first > 8191 || last > 8191 {
raise Invalid("PIZ bitmap range", 0)
}
let lut = [0]
for i = first; i <= last; i = i + 1 {
let entry = r.u8()
for bit = 0; bit < 8; bit = bit + 1 {
let word = i * 8 + bit
if word != 0 && (entry & (1 << bit)) != 0 {
lut.push(word)
}
}
}
let length = r.i32()
if length != r.end - r.pos {
raise Invalid("PIZ Huffman payload size", r.pos)
}
let words = huffman_decode(r.take(length), expected / 2)
piz_wavelet(words, channels, w, lut.length() - 1, true)
let positions = piz_positions(channels, w)
let out = sink()
for _row = 0; _row < w.height(); _row = _row + 1 {
for i = 0; i < channels.length(); i = i + 1 {
let row_words = w.width() * (channels[i].kind.bytes() / 2)
for j = 0; j < row_words; j = j + 1 {
let rank = words[positions[i] + j]
if rank >= lut.length() {
raise Invalid("PIZ coefficient outside range map", 0)
}
out.uint(lut[rank].to_uint64(), 2)
}
positions[i] += row_words
}
}
out.finish()
}
///|
fn encode_block(
raw : Bytes,
codec : Compression,
channels : Array[Channel],
window : Window,
) -> Bytes raise ExrError {
if codec == Piz {
piz_pack(raw, channels, window)
} else {
pack(raw, codec)
}
}
///|
fn decode_block(
bytes : Bytes,
codec : Compression,
channels : Array[Channel],
window : Window,
expected : Int32,
) -> Bytes raise ExrError {
if codec == Piz && bytes.length() < expected {
piz_unpack(bytes, channels, window, expected)
} else {
unpack(bytes, codec, expected)
}
}