Skip to content

Commit 3caf54e

Browse files
committed
Simplify flag parsing
1 parent 29d0051 commit 3caf54e

File tree

1 file changed

+33
-51
lines changed

1 file changed

+33
-51
lines changed

crates/iceberg/src/puffin/metadata.rs

Lines changed: 33 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
use std::collections::{HashMap, HashSet};
1919

2020
use bytes::Bytes;
21-
use once_cell::sync::Lazy;
2221
use serde::{Deserialize, Serialize};
2322

2423
use crate::io::{FileRead, InputFile};
@@ -56,44 +55,37 @@ pub(crate) struct BlobMetadata {
5655
pub(crate) properties: HashMap<String, String>,
5756
}
5857

59-
#[derive(Clone, PartialEq, Eq, Hash, Debug)]
58+
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
6059
pub(crate) enum Flag {
61-
FooterPayloadCompressed,
60+
FooterPayloadCompressed = 0
6261
}
6362

64-
#[derive(PartialEq, Eq, Hash)]
65-
pub(crate) struct ByteNumber(pub u8);
66-
67-
#[derive(PartialEq, Eq, Hash)]
68-
pub(crate) struct BitNumber(pub u8);
69-
70-
static FLAGS_BY_BYTE_AND_BIT: Lazy<HashMap<(ByteNumber, BitNumber), Flag>> = Lazy::new(|| {
71-
let mut m = HashMap::new();
72-
m.insert(
73-
(
74-
Flag::FooterPayloadCompressed.byte_number(),
75-
Flag::FooterPayloadCompressed.bit_number(),
76-
),
77-
Flag::FooterPayloadCompressed,
78-
);
79-
m
80-
});
81-
8263
impl Flag {
83-
pub(crate) fn byte_number(&self) -> ByteNumber {
84-
match self {
85-
Flag::FooterPayloadCompressed => ByteNumber(0),
86-
}
64+
pub(crate) fn byte_idx(self) -> u8 {
65+
(self as u8) / 8
66+
}
67+
68+
pub(crate) fn bit_idx(self) -> u8 {
69+
(self as u8) % 8
8770
}
8871

89-
pub(crate) fn bit_number(&self) -> BitNumber {
90-
match self {
91-
Flag::FooterPayloadCompressed => BitNumber(0),
92-
}
72+
fn matches(self, byte_idx: &u8, bit_idx: &u8) -> bool {
73+
&self.byte_idx() == byte_idx &&
74+
&self.bit_idx() == bit_idx
9375
}
9476

95-
fn from(byte_and_bit: &(ByteNumber, BitNumber)) -> Option<Flag> {
96-
FLAGS_BY_BYTE_AND_BIT.get(byte_and_bit).cloned()
77+
fn from(byte_idx: &u8, bit_idx: &u8) -> Result<Flag> {
78+
if Flag::FooterPayloadCompressed.matches(byte_idx, bit_idx) {
79+
Ok(Flag::FooterPayloadCompressed)
80+
} else {
81+
Err(Error::new(
82+
ErrorKind::DataInvalid,
83+
format!(
84+
"Unknown flag byte {} and bit {} combination",
85+
byte_idx, bit_idx
86+
),
87+
))
88+
}
9789
}
9890
}
9991

@@ -178,35 +170,25 @@ impl FileMetadata {
178170

179171
fn decode_flags(footer_bytes: &[u8]) -> Result<HashSet<Flag>> {
180172
let mut flags = HashSet::new();
181-
for byte_number in 0..FileMetadata::FOOTER_STRUCT_FLAGS_LENGTH {
173+
174+
for byte_idx in 0..FileMetadata::FOOTER_STRUCT_FLAGS_LENGTH {
182175
let byte_offset = footer_bytes.len()
183176
- usize::from(FileMetadata::MAGIC_LENGTH)
184177
- usize::from(FileMetadata::FOOTER_STRUCT_FLAGS_LENGTH)
185-
+ usize::from(byte_number);
178+
+ usize::from(byte_idx);
186179

187-
let mut flag_byte = *footer_bytes.get(byte_offset).ok_or_else(|| {
180+
let flag_byte = *footer_bytes.get(byte_offset).ok_or_else(|| {
188181
Error::new(ErrorKind::DataInvalid, "Index range is out of bounds.")
189182
})?;
190-
let mut bit_number = 0;
191-
while flag_byte != 0 {
192-
if flag_byte & 0x1 != 0 {
193-
match Flag::from(&(ByteNumber(byte_number), BitNumber(bit_number))) {
194-
Some(flag) => flags.insert(flag),
195-
None => {
196-
return Err(Error::new(
197-
ErrorKind::DataInvalid,
198-
format!(
199-
"Unknown flag byte {} and bit {} combination",
200-
byte_number, bit_number
201-
),
202-
))
203-
}
204-
};
183+
184+
for bit_idx in 0..8 {
185+
if ((flag_byte >> bit_idx) & 1) != 0 {
186+
let flag = Flag::from(&byte_idx, &bit_idx)?;
187+
flags.insert(flag);
205188
}
206-
flag_byte >>= 1;
207-
bit_number += 1;
208189
}
209190
}
191+
210192
Ok(flags)
211193
}
212194

0 commit comments

Comments
 (0)