prolly/prolly/proximity/accelerator/
catalog.rs

1use super::composite::{config_fingerprint as composite_fingerprint, CompositeAccelerator};
2use super::hnsw::storage::config_fingerprint as hnsw_fingerprint;
3use super::pq::config_fingerprint as pq_fingerprint;
4use super::{AcceleratorSet, HnswIndex, ProductQuantizer};
5use crate::prolly::cid::Cid;
6use crate::prolly::content_graph::{ContentObjectKind, TypedContentRoot};
7use crate::prolly::error::Error;
8use crate::prolly::proximity::storage::codec::{put_cid, put_varint, Reader, MAX_OBJECT_ENTRIES};
9use crate::prolly::proximity::ProximityTree;
10use crate::prolly::store::{NodePublication, PublicationOrigin, Store};
11
12const MAGIC: &[u8; 4] = b"PACL";
13const VERSION: u8 = 1;
14
15#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
16pub enum CatalogAcceleratorKind {
17    Hnsw,
18    ProductQuantized,
19    Composite,
20}
21
22impl CatalogAcceleratorKind {
23    pub(crate) const fn id(self) -> u8 {
24        match self {
25            Self::Hnsw => 1,
26            Self::ProductQuantized => 2,
27            Self::Composite => 3,
28        }
29    }
30
31    fn from_id(id: u8) -> Result<Self, Error> {
32        match id {
33            1 => Ok(Self::Hnsw),
34            2 => Ok(Self::ProductQuantized),
35            3 => Ok(Self::Composite),
36            _ => Err(invalid("unknown catalog accelerator kind")),
37        }
38    }
39}
40
41#[derive(Clone, Debug, PartialEq, Eq)]
42pub struct AcceleratorCatalogEntry {
43    pub kind: CatalogAcceleratorKind,
44    pub configuration_fingerprint: Cid,
45    pub manifest: Cid,
46}
47
48pub struct AcceleratorCatalog<S: Store> {
49    manifest: Cid,
50    source: Cid,
51    entries: Vec<AcceleratorCatalogEntry>,
52    accelerators: AcceleratorSet<S>,
53}
54
55impl<S> AcceleratorCatalog<S>
56where
57    S: Store + Clone + Send + Sync,
58    S::Error: Send + Sync,
59{
60    pub fn build(
61        store: S,
62        source: &ProximityTree,
63        accelerators: AcceleratorSet<S>,
64    ) -> Result<Self, Error> {
65        let entries = entries_from_set(&accelerators);
66        if entries.is_empty() {
67            return Err(invalid("accelerator catalog must not be empty"));
68        }
69        let object = Manifest {
70            source: source.descriptor.clone(),
71            entries: entries.clone(),
72        };
73        let bytes = object.encode()?;
74        let manifest = Cid::from_bytes(&bytes);
75        put_content(&store, &manifest, &bytes)?;
76        Ok(Self {
77            manifest,
78            source: source.descriptor.clone(),
79            entries,
80            accelerators,
81        })
82    }
83
84    pub fn load(store: S, manifest: Cid, source: &ProximityTree) -> Result<Self, Error> {
85        let bytes = load_content(&store, &manifest)?;
86        let object = Manifest::decode(&bytes)?;
87        if object.source != source.descriptor {
88            return Err(invalid("catalog is bound to a different source snapshot"));
89        }
90        let mut accelerators = AcceleratorSet::empty();
91        for entry in &object.entries {
92            accelerators = match entry.kind {
93                CatalogAcceleratorKind::Hnsw => {
94                    let index = HnswIndex::load(store.clone(), entry.manifest.clone())?;
95                    if hnsw_fingerprint(index.config()) != entry.configuration_fingerprint {
96                        return Err(invalid("catalog HNSW fingerprint mismatch"));
97                    }
98                    accelerators.with_hnsw(source, index)?
99                }
100                CatalogAcceleratorKind::ProductQuantized => {
101                    let index = ProductQuantizer::load(store.clone(), entry.manifest.clone())?;
102                    if pq_fingerprint(index.config()) != entry.configuration_fingerprint {
103                        return Err(invalid("catalog PQ fingerprint mismatch"));
104                    }
105                    accelerators.with_pq(source, index)?
106                }
107                CatalogAcceleratorKind::Composite => {
108                    let index = CompositeAccelerator::load(store.clone(), entry.manifest.clone())?;
109                    if composite_fingerprint(index.config()) != entry.configuration_fingerprint {
110                        return Err(invalid("catalog composite fingerprint mismatch"));
111                    }
112                    accelerators.with_composite(source, index)?
113                }
114            };
115        }
116        Ok(Self {
117            manifest,
118            source: object.source,
119            entries: object.entries,
120            accelerators,
121        })
122    }
123
124    pub fn manifest_cid(&self) -> &Cid {
125        &self.manifest
126    }
127    pub fn typed_root(&self) -> TypedContentRoot {
128        TypedContentRoot::new(ContentObjectKind::AcceleratorCatalog, self.manifest.clone())
129    }
130    pub fn source_descriptor(&self) -> &Cid {
131        &self.source
132    }
133    pub fn entries(&self) -> &[AcceleratorCatalogEntry] {
134        &self.entries
135    }
136    pub fn accelerators(&self) -> &AcceleratorSet<S> {
137        &self.accelerators
138    }
139    pub fn into_accelerators(self) -> AcceleratorSet<S> {
140        self.accelerators
141    }
142}
143
144fn entries_from_set<S>(set: &AcceleratorSet<S>) -> Vec<AcceleratorCatalogEntry>
145where
146    S: Store + Clone + Send + Sync,
147    S::Error: Send + Sync,
148{
149    let mut entries = Vec::new();
150    if let Some(index) = set.hnsw() {
151        entries.push(AcceleratorCatalogEntry {
152            kind: CatalogAcceleratorKind::Hnsw,
153            configuration_fingerprint: hnsw_fingerprint(index.config()),
154            manifest: index.manifest_cid().clone(),
155        });
156    }
157    if let Some(index) = set.pq() {
158        entries.push(AcceleratorCatalogEntry {
159            kind: CatalogAcceleratorKind::ProductQuantized,
160            configuration_fingerprint: pq_fingerprint(index.config()),
161            manifest: index.manifest_cid().clone(),
162        });
163    }
164    if let Some(index) = set.composite() {
165        entries.push(AcceleratorCatalogEntry {
166            kind: CatalogAcceleratorKind::Composite,
167            configuration_fingerprint: composite_fingerprint(index.config()),
168            manifest: index.manifest_cid().clone(),
169        });
170    }
171    entries.sort_by(compare_entry);
172    entries
173}
174
175#[derive(Clone)]
176pub(crate) struct Manifest {
177    pub(crate) source: Cid,
178    pub(crate) entries: Vec<AcceleratorCatalogEntry>,
179}
180
181impl Manifest {
182    pub(crate) fn encode(&self) -> Result<Vec<u8>, Error> {
183        self.validate()?;
184        let mut bytes = Vec::new();
185        bytes.extend_from_slice(MAGIC);
186        bytes.push(VERSION);
187        put_cid(&self.source, &mut bytes);
188        put_varint(self.entries.len() as u64, &mut bytes);
189        for entry in &self.entries {
190            bytes.push(entry.kind.id());
191            put_cid(&entry.configuration_fingerprint, &mut bytes);
192            put_cid(&entry.manifest, &mut bytes);
193        }
194        Ok(bytes)
195    }
196
197    pub(crate) fn decode(bytes: &[u8]) -> Result<Self, Error> {
198        let mut reader = Reader::new(bytes, "accelerator catalog");
199        reader.exact(MAGIC)?;
200        if reader.u8()? != VERSION {
201            return Err(reader.invalid("unsupported accelerator catalog version"));
202        }
203        let source = reader.cid()?;
204        let count = reader.bounded_usize(MAX_OBJECT_ENTRIES)?;
205        let mut entries = Vec::with_capacity(count);
206        for _ in 0..count {
207            entries.push(AcceleratorCatalogEntry {
208                kind: CatalogAcceleratorKind::from_id(reader.u8()?)?,
209                configuration_fingerprint: reader.cid()?,
210                manifest: reader.cid()?,
211            });
212        }
213        reader.finish()?;
214        let object = Self { source, entries };
215        object.validate()?;
216        Ok(object)
217    }
218
219    fn validate(&self) -> Result<(), Error> {
220        if self.entries.is_empty()
221            || self.entries.windows(2).any(|pair| {
222                compare_entry(&pair[0], &pair[1]) != std::cmp::Ordering::Less
223                    || pair[0].kind == pair[1].kind
224            })
225        {
226            return Err(invalid(
227                "catalog entries must be sorted, unique, and contain one entry per kind",
228            ));
229        }
230        Ok(())
231    }
232}
233
234fn compare_entry(
235    left: &AcceleratorCatalogEntry,
236    right: &AcceleratorCatalogEntry,
237) -> std::cmp::Ordering {
238    left.kind
239        .cmp(&right.kind)
240        .then_with(|| {
241            left.configuration_fingerprint
242                .as_bytes()
243                .cmp(right.configuration_fingerprint.as_bytes())
244        })
245        .then_with(|| left.manifest.as_bytes().cmp(right.manifest.as_bytes()))
246}
247
248fn load_content<S: Store>(store: &S, cid: &Cid) -> Result<Vec<u8>, Error> {
249    let bytes = store
250        .get(cid.as_bytes())
251        .map_err(|error| Error::Store(Box::new(error)))?
252        .ok_or_else(|| Error::NotFound(cid.clone()))?;
253    let actual = Cid::from_bytes(&bytes);
254    if actual != *cid {
255        return Err(Error::CidMismatch {
256            expected: cid.clone(),
257            actual,
258        });
259    }
260    Ok(bytes)
261}
262
263fn put_content<S: Store>(store: &S, cid: &Cid, bytes: &[u8]) -> Result<(), Error> {
264    if let Some(existing) = store
265        .get(cid.as_bytes())
266        .map_err(|error| Error::Store(Box::new(error)))?
267    {
268        let actual = Cid::from_bytes(&existing);
269        if actual != *cid {
270            return Err(Error::CidMismatch {
271                expected: cid.clone(),
272                actual,
273            });
274        }
275        return Ok(());
276    }
277    let entries = [(cid.as_bytes(), bytes)];
278    store
279        .publish_nodes(NodePublication::new(
280            &entries,
281            PublicationOrigin::Maintenance,
282        ))
283        .map_err(|error| Error::Store(Box::new(error)))
284}
285
286fn invalid(reason: impl Into<String>) -> Error {
287    Error::InvalidProximityObject {
288        kind: "accelerator catalog",
289        reason: reason.into(),
290    }
291}
prolly/prolly/proximity/accelerator/catalog.rs

prolly/prolly/proximity/accelerator/
catalog.rs
