// Adapted from gitoxide's gix-bitmap (MIT OR Apache-2.0).
use knot_git::GitError;
const RUN_LEN_MAX: u64 = (1 << 32) - 1;
const LITERALS_MAX: u64 = (1 << 31) - 1;
#[derive(Clone)]
pub(crate) struct Ewah {
num_bits: u32,
words: Vec<u64>,
rlw: u32,
}
impl Ewah {
pub(crate) fn compress(input: &[u64], num_bits: usize) -> Result<Self, GitError> {
let input = &input[..num_bits.div_ceil(64)];
let num_bits = u32::try_from(num_bits)
.map_err(|_| GitError::Backend("ewah bit count exceeds u32".to_string()))?;
let full = num_bits as usize / 64;
let clean = |at: usize| at < full && (input[at] == 0 || input[at] == !0);
let mut words = Vec::new();
let mut rlw;
let mut at = 0;
loop {
rlw = words.len();
words.push(0u64);
let run_bit = clean(at) && input[at] == !0;
let mut run = 0u64;
while clean(at) && (input[at] == !0) == run_bit && run < RUN_LEN_MAX {
run += 1;
at += 1;
}
let mut literals = 0u64;
while at < input.len() && !clean(at) && literals < LITERALS_MAX {
words.push(input[at]);
literals += 1;
at += 1;
}
words[rlw] = u64::from(run_bit) | run << 1 | literals << 33;
if at == input.len() {
break;
}
}
Ok(Self {
num_bits,
words,
rlw: rlw as u32,
})
}
pub(crate) fn decode(data: &[u8]) -> Result<(Self, &[u8]), GitError> {
let (num_bits, data) = take_u32(data)?;
let (len, data) = take_u32(data)?;
let (body, data) = (len as usize)
.checked_mul(8)
.and_then(|size| data.split_at_checked(size))
.ok_or_else(truncated)?;
let words = body
.as_chunks::<8>()
.0
.iter()
.map(|word| u64::from_be_bytes(*word))
.collect();
let (rlw, data) = take_u32(data)?;
Ok((
Self {
num_bits,
words,
rlw,
},
data,
))
}
pub(crate) fn write_to(&self, out: &mut Vec<u8>) {
out.extend_from_slice(&self.num_bits.to_be_bytes());
out.extend_from_slice(&(self.words.len() as u32).to_be_bytes());
self.words
.iter()
.for_each(|word| out.extend_from_slice(&word.to_be_bytes()));
out.extend_from_slice(&self.rlw.to_be_bytes());
}
pub(crate) fn word_count(&self) -> usize {
self.words.len()
}
pub(crate) fn for_each_word(
&self,
mut f: impl FnMut(usize, u64) -> Result<(), GitError>,
) -> Result<(), GitError> {
let limit = (self.num_bits as usize).div_ceil(64);
let malformed = || GitError::Backend("malformed ewah bitmap".to_string());
let mut at = 0usize;
let mut words = self.words.iter();
while let Some(marker) = words.next() {
let run = ((marker >> 1) & RUN_LEN_MAX) as usize;
let end = at
.checked_add(run)
.filter(|end| *end <= limit)
.ok_or_else(malformed)?;
if marker & 1 == 1 {
(at..end).try_for_each(|index| f(index, !0))?;
}
at = end;
for _ in 0..marker >> 33 {
let word = words.next().filter(|_| at < limit).ok_or_else(malformed)?;
f(at, *word)?;
at += 1;
}
}
Ok(())
}
}
fn take_u32(data: &[u8]) -> Result<(u32, &[u8]), GitError> {
data.split_at_checked(4)
.map(|(num, rest)| (u32::from_be_bytes(num.try_into().unwrap()), rest))
.ok_or_else(truncated)
}
fn truncated() -> GitError {
GitError::Backend("truncated ewah bitmap".to_string())
}
#[cfg(test)]
mod tests {
use super::*;
fn encode(num_bits: u32, words: &[u64]) -> Vec<u8> {
let mut out = Vec::new();
Ewah {
num_bits,
words: words.to_vec(),
rlw: 0,
}
.write_to(&mut out);
out
}
fn walk(bytes: &[u8]) -> Result<Vec<(usize, u64)>, GitError> {
let (vector, _) = Ewah::decode(bytes)?;
let mut seen = Vec::new();
vector.for_each_word(|index, word| {
seen.push((index, word));
Ok(())
})?;
Ok(seen)
}
#[test]
fn rejects_truncation_and_words_past_the_bit_count() {
assert_eq!(
walk(&encode(128, &[1 | 1 << 1 | 1 << 33, 5])).unwrap(),
vec![(0, !0), (1, 5)]
);
let whole = encode(64, &[1 << 33, 5]);
assert!(walk(&whole[..whole.len() - 1]).is_err());
assert!(walk(&encode(64, &[2 << 1])).is_err(), "run past num_bits");
assert!(
walk(&encode(64, &[2 << 33, 1, 2])).is_err(),
"literal past num_bits"
);
assert!(
walk(&encode(64, &[1 << 33])).is_err(),
"missing literal word"
);
}
}