Skip to main content

heddle_pack/store/pack/
pack_index.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Pack index for fast object lookup within packfiles.
3
4use bytes::Bytes;
5
6use crate::store::{
7    Result,
8    pack::{
9        PackObjectId,
10        versioned_header::{HeaderChecksum, VersionedHeader},
11    },
12};
13
14pub(super) const INDEX_MAGIC: &[u8; 4] = b"LMI\0";
15pub(super) const INDEX_VERSION: u32 = 4;
16pub(super) const INDEX_ENTRY_LEN: usize = 32 + 8;
17const STATE_ID_OFFSET_TAG: u64 = 1 << 63;
18const ANNOTATED_TAG_OFFSET_TAG: u64 = 1 << 62;
19const PACK_OFFSET_MASK: u64 = !(STATE_ID_OFFSET_TAG | ANNOTATED_TAG_OFFSET_TAG);
20
21/// Entry in the pack index.
22#[derive(Debug, Clone, Copy)]
23pub struct IndexEntry {
24    pub id: PackObjectId,
25    pub offset: u64,
26}
27
28/// Pack index for fast object lookup.
29#[derive(Debug)]
30pub struct PackIndex {
31    entries: Vec<IndexEntry>,
32    encoded: Option<EncodedIndex>,
33}
34
35#[derive(Debug)]
36struct EncodedIndex {
37    data: Bytes,
38    entries_start: usize,
39    count: usize,
40}
41
42impl PackIndex {
43    /// Create a new empty index.
44    pub fn new() -> Self {
45        Self {
46            entries: Vec::new(),
47            encoded: None,
48        }
49    }
50
51    /// Add an entry.
52    pub fn add(&mut self, id: PackObjectId, offset: u64) {
53        debug_assert!(self.encoded.is_none());
54        self.entries.push(IndexEntry { id, offset });
55    }
56
57    /// Sort entries by hash for binary search.
58    pub fn sort(&mut self) {
59        debug_assert!(self.encoded.is_none());
60        self.entries.sort_by_key(|e| e.id);
61    }
62
63    /// Find an entry by hash.
64    pub fn find(&self, id: &PackObjectId) -> Result<Option<u64>> {
65        let Some(encoded) = &self.encoded else {
66            return Ok(self
67                .entries
68                .binary_search_by_key(id, |entry| entry.id)
69                .ok()
70                .map(|index| self.entries[index].offset));
71        };
72        let mut low = 0;
73        let mut high = encoded.count;
74        while low < high {
75            let middle = low + (high - low) / 2;
76            let entry = encoded.entry(middle)?;
77            match entry.id.cmp(id) {
78                std::cmp::Ordering::Less => low = middle + 1,
79                std::cmp::Ordering::Greater => high = middle,
80                std::cmp::Ordering::Equal => return Ok(Some(entry.offset)),
81            }
82        }
83        Ok(None)
84    }
85
86    /// Serialize to bytes.
87    pub fn to_bytes(&self) -> Vec<u8> {
88        if let Some(encoded) = &self.encoded {
89            return encoded.data.to_vec();
90        }
91        let mut result = Vec::new();
92        index_header().write_vec(&mut result, self.entries.len() as u64);
93        for entry in &self.entries {
94            result.extend_from_slice(&encode_index_entry(entry.id, entry.offset));
95        }
96        result
97    }
98
99    /// Deserialize from bytes.
100    pub fn from_bytes(data: &[u8]) -> Result<Self> {
101        Self::from_owned_bytes(Bytes::copy_from_slice(data))
102    }
103
104    pub fn from_owned_bytes(data: Bytes) -> Result<Self> {
105        verify_index_version(&data)?;
106        let header = index_header().verify(&data)?;
107        let count = header.count;
108        let max_entries = ((data.len() - header.header_len) / INDEX_ENTRY_LEN) as u64;
109        if count > max_entries {
110            return Err(crate::store::StoreError::InvalidObject(format!(
111                "Index entry count {} exceeds available data capacity {}",
112                count, max_entries
113            )));
114        }
115        let count = usize::try_from(count).map_err(|_| {
116            crate::store::StoreError::InvalidObject(
117                "Index entry count exceeds platform limits".to_string(),
118            )
119        })?;
120        let index = Self {
121            entries: Vec::new(),
122            encoded: Some(EncodedIndex {
123                data,
124                entries_start: header.header_len,
125                count,
126            }),
127        };
128        if header.header_len + count * INDEX_ENTRY_LEN
129            != index
130                .encoded
131                .as_ref()
132                .map_or(0, |encoded| encoded.data.len())
133        {
134            return Err(crate::store::StoreError::InvalidObject(
135                "index has trailing bytes".into(),
136            ));
137        }
138        let mut previous = None;
139        for entry in index.iter() {
140            let entry = entry?;
141            if previous.is_some_and(|id| id > entry.id) {
142                return Err(crate::store::StoreError::InvalidObject(
143                    "index identities are not ordered".into(),
144                ));
145            }
146            previous = Some(entry.id);
147        }
148        Ok(index)
149    }
150}
151
152impl EncodedIndex {
153    fn entry(&self, index: usize) -> Result<IndexEntry> {
154        let start = self.entries_start + index * INDEX_ENTRY_LEN;
155        let end = start + INDEX_ENTRY_LEN;
156        let bytes = self.data.get(start..end).ok_or_else(|| {
157            crate::store::StoreError::InvalidObject("Index data truncated".to_string())
158        })?;
159        decode_index_entry(bytes)
160    }
161}
162
163pub(super) fn encode_index_entry(id: PackObjectId, offset: u64) -> [u8; INDEX_ENTRY_LEN] {
164    assert!(
165        offset <= PACK_OFFSET_MASK,
166        "pack index offset exceeds the 62-bit format limit"
167    );
168    let mut bytes = [0u8; INDEX_ENTRY_LEN];
169    let tagged_offset = match id {
170        PackObjectId::Hash(hash) => {
171            bytes[..32].copy_from_slice(hash.as_bytes());
172            offset
173        }
174        PackObjectId::StateId(state_id) => {
175            bytes[..32].copy_from_slice(state_id.as_bytes());
176            offset | STATE_ID_OFFSET_TAG
177        }
178        PackObjectId::AnnotatedTag(hash) => {
179            bytes[..32].copy_from_slice(hash.as_bytes());
180            offset | ANNOTATED_TAG_OFFSET_TAG
181        }
182    };
183    bytes[32..].copy_from_slice(&tagged_offset.to_be_bytes());
184    bytes
185}
186
187fn decode_index_entry(bytes: &[u8]) -> Result<IndexEntry> {
188    let raw_id: [u8; 32] = bytes[..32].try_into().map_err(|_| {
189        crate::store::StoreError::InvalidObject("Invalid index id length".to_string())
190    })?;
191    let tagged_offset = u64::from_be_bytes(bytes[32..].try_into().map_err(|_| {
192        crate::store::StoreError::InvalidObject("Invalid offset length".to_string())
193    })?);
194    let id = if tagged_offset & STATE_ID_OFFSET_TAG != 0 {
195        PackObjectId::StateId(crate::object::StateId::from_bytes(raw_id))
196    } else if tagged_offset & ANNOTATED_TAG_OFFSET_TAG != 0 {
197        PackObjectId::AnnotatedTag(crate::object::ContentHash::from_bytes(raw_id))
198    } else {
199        PackObjectId::Hash(crate::object::ContentHash::from_bytes(raw_id))
200    };
201    Ok(IndexEntry {
202        id,
203        offset: tagged_offset & PACK_OFFSET_MASK,
204    })
205}
206
207impl PackIndex {
208    /// Return all decoded index entries.
209    pub(super) fn entries(&self) -> Result<Vec<IndexEntry>> {
210        if let Some(encoded) = &self.encoded {
211            return (0..encoded.count)
212                .map(|index| encoded.entry(index))
213                .collect();
214        }
215        Ok(self.entries.clone())
216    }
217
218    /// Return all ids in this index.
219    pub fn ids(&self) -> Result<Vec<PackObjectId>> {
220        Ok(self.entries()?.into_iter().map(|entry| entry.id).collect())
221    }
222
223    pub fn len(&self) -> usize {
224        self.encoded
225            .as_ref()
226            .map_or(self.entries.len(), |index| index.count)
227    }
228    pub fn is_empty(&self) -> bool {
229        self.len() == 0
230    }
231    pub(super) fn entry(&self, index: usize) -> Result<IndexEntry> {
232        match &self.encoded {
233            Some(encoded) => encoded.entry(index),
234            None => self.entries.get(index).copied().ok_or_else(|| {
235                crate::store::StoreError::InvalidObject("index entry out of range".into())
236            }),
237        }
238    }
239    pub(super) fn iter(&self) -> impl Iterator<Item = Result<IndexEntry>> + '_ {
240        (0..self.len()).map(|index| self.entry(index))
241    }
242}
243
244impl Default for PackIndex {
245    fn default() -> Self {
246        Self::new()
247    }
248}
249
250fn verify_index_version(data: &[u8]) -> Result<()> {
251    if data.len() < 8 || &data[..4] != INDEX_MAGIC {
252        index_header().verify_layout(data)?;
253        unreachable!("invalid index header must have returned an error")
254    }
255    let version = u32::from_be_bytes(data[4..8].try_into().map_err(|_| {
256        crate::store::StoreError::InvalidObject("Index version field is truncated".to_string())
257    })?);
258    if version == INDEX_VERSION {
259        Ok(())
260    } else if version > INDEX_VERSION {
261        Err(crate::store::StoreError::InvalidObject(format!(
262            "pack index uses format version {version}, but this binary supports {INDEX_VERSION}; upgrade heddle"
263        )))
264    } else {
265        Err(crate::store::StoreError::InvalidObject(format!(
266            "pack index uses unsupported format version {version}; recreate the repository or re-adopt its Git history with this Heddle version"
267        )))
268    }
269}
270
271pub(super) fn index_header() -> VersionedHeader {
272    VersionedHeader {
273        magic: INDEX_MAGIC,
274        version: INDEX_VERSION,
275        checksum: HeaderChecksum::None,
276        too_short: "Index too short",
277        invalid_magic: "Invalid index magic",
278        unsupported_version: "Unsupported index version",
279        checksum_mismatch: "",
280    }
281}