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

heddle_pack/store/pack/
pack_index.rs
