molframe_core/provider/
source.rs

1//! Declarative adapters from immutable structure snapshots to provider chunks.
2
3use super::{
4    BondChunk, ChunkDescriptor, ChunkId, DatasetDescriptor, DatasetId, FrameChunk, LogicalRow,
5    PayloadKind, PropertyChunk, ProviderError, StructureChunk,
6};
7use crate::chunk::TARGET_CHUNK_ATOMS;
8use crate::structure::Structure;
9use std::sync::Arc;
10
11/// Default maximum number of bonds independently requested in one payload.
12pub const TARGET_CHUNK_BONDS: u32 = 65_536;
13
14/// Zero-copy provider for regularly paged chemical bond topology.
15#[derive(Clone, Debug)]
16pub struct BondChunkProvider {
17    structure: Structure,
18    dataset: DatasetDescriptor,
19    atom_dataset: DatasetId,
20    atom_logical_start: LogicalRow,
21}
22
23impl BondChunkProvider {
24    /// Creates a provider using the default bounded bond payload size.
25    ///
26    /// # Errors
27    ///
28    /// Rejects unavailable topology and identity or chunk-range overflow.
29    pub fn new(
30        dataset: DatasetId,
31        first_chunk: ChunkId,
32        atom_dataset: DatasetId,
33        atom_logical_start: LogicalRow,
34        structure: Structure,
35    ) -> Result<Self, ProviderError> {
36        Self::with_rows_per_chunk(
37            dataset,
38            first_chunk,
39            atom_dataset,
40            atom_logical_start,
41            structure,
42            TARGET_CHUNK_BONDS,
43        )
44    }
45
46    /// Creates a provider with an explicit bounded payload row count.
47    ///
48    /// # Errors
49    ///
50    /// Rejects zero chunk rows, unavailable topology, or identity overflow.
51    pub fn with_rows_per_chunk(
52        dataset: DatasetId,
53        first_chunk: ChunkId,
54        atom_dataset: DatasetId,
55        atom_logical_start: LogicalRow,
56        structure: Structure,
57        rows_per_chunk: u32,
58    ) -> Result<Self, ProviderError> {
59        if !structure.data().bonds.is_available() {
60            return Err(ProviderError::BondTopologyUnavailable);
61        }
62        validate_atom_endpoint_range(atom_dataset, atom_logical_start, structure.atom_count())?;
63        let logical_rows = u64::try_from(structure.data().bonds.len()).map_err(|_| {
64            ProviderError::IdentityOverflow {
65                dataset,
66                field: "bond table rows",
67            }
68        })?;
69        let descriptor = DatasetDescriptor::regular(
70            dataset,
71            PayloadKind::BondTopology,
72            logical_rows,
73            rows_per_chunk,
74            first_chunk,
75        )?;
76        Ok(Self {
77            structure,
78            dataset: descriptor,
79            atom_dataset,
80            atom_logical_start,
81        })
82    }
83
84    /// Compact dataset metadata; no per-chunk index is retained.
85    #[must_use]
86    pub const fn dataset(&self) -> DatasetDescriptor {
87        self.dataset
88    }
89
90    /// Produces one shared topology payload in O(1), without a table scan.
91    ///
92    /// # Errors
93    ///
94    /// Rejects chunk identities outside this provider or source index overflow.
95    pub fn chunk(&self, chunk: ChunkId) -> Result<BondChunk, ProviderError> {
96        let descriptor = self.dataset.regular_chunk(chunk)?;
97        let start = u32::try_from(descriptor.logical_start().get()).map_err(|_| {
98            ProviderError::IdentityOverflow {
99                dataset: self.dataset.id(),
100                field: "bond source row",
101            }
102        })?;
103        let end = start
104            .checked_add(descriptor.rows())
105            .ok_or(ProviderError::IdentityOverflow {
106                dataset: self.dataset.id(),
107                field: "bond source range",
108            })?;
109        BondChunk::shared(
110            descriptor,
111            self.structure.clone(),
112            start..end,
113            self.atom_dataset,
114            self.atom_logical_start,
115        )
116    }
117}
118
119fn storage_descriptor(
120    dataset: &DatasetDescriptor,
121    structure: &Structure,
122    chunk: ChunkId,
123) -> Result<(usize, ChunkDescriptor), ProviderError> {
124    let ordinal = dataset.chunk_ordinal(chunk)?;
125    let index = usize::try_from(ordinal).map_err(|_| ProviderError::IdentityOverflow {
126        dataset: dataset.id(),
127        field: "storage chunk index",
128    })?;
129    let storage = structure
130        .data()
131        .chunks
132        .get(index)
133        .ok_or(ProviderError::StorageChunkUnavailable { index })?;
134    let descriptor = ChunkDescriptor::new(
135        dataset.id(),
136        chunk,
137        LogicalRow::new(u64::from(storage.atoms().start)),
138        storage.len(),
139    )?;
140    Ok((index, descriptor))
141}
142
143pub(super) fn validate_atom_endpoint_range(
144    dataset: DatasetId,
145    first: LogicalRow,
146    rows: u32,
147) -> Result<(), ProviderError> {
148    if rows == 0 {
149        return Ok(());
150    }
151    first
152        .get()
153        .checked_add(u64::from(rows - 1))
154        .map(|_| ())
155        .ok_or(ProviderError::IdentityOverflow {
156            dataset,
157            field: "atom endpoint range",
158        })
159}
160
161/// Zero-copy provider for structure chunks already present in one snapshot.
162#[derive(Clone, Debug)]
163pub struct StructureChunkProvider {
164    structure: Structure,
165    dataset: DatasetDescriptor,
166}
167
168impl StructureChunkProvider {
169    /// Creates a provider whose chunk namespace begins at `first_chunk`.
170    ///
171    /// # Errors
172    ///
173    /// Returns a typed error when chunk count or identity arithmetic overflows.
174    pub fn new(
175        dataset: DatasetId,
176        first_chunk: ChunkId,
177        structure: Structure,
178    ) -> Result<Self, ProviderError> {
179        let chunk_count = u64::try_from(structure.data().chunks.len()).map_err(|_| {
180            ProviderError::IdentityOverflow {
181                dataset,
182                field: "structure chunk count",
183            }
184        })?;
185        let descriptor = DatasetDescriptor::source_defined(
186            dataset,
187            PayloadKind::Structure,
188            u64::from(structure.atom_count()),
189            chunk_count,
190            first_chunk,
191            TARGET_CHUNK_ATOMS,
192        )?;
193        Ok(Self {
194            structure,
195            dataset: descriptor,
196        })
197    }
198
199    /// Compact dataset metadata.
200    #[must_use]
201    pub const fn dataset(&self) -> DatasetDescriptor {
202        self.dataset
203    }
204
205    /// Produces one shared structure payload in constant time.
206    ///
207    /// # Errors
208    ///
209    /// Rejects chunk identities outside this provider.
210    pub fn chunk(&self, chunk: ChunkId) -> Result<StructureChunk, ProviderError> {
211        let (index, descriptor) = storage_descriptor(&self.dataset, &self.structure, chunk)?;
212        StructureChunk::shared(descriptor, self.structure.clone(), index)
213    }
214}
215
216/// Zero-copy provider for one typed atom property.
217#[derive(Clone, Debug)]
218pub struct PropertyChunkProvider {
219    structure: Structure,
220    name: Arc<str>,
221    dataset: DatasetDescriptor,
222}
223
224impl PropertyChunkProvider {
225    /// Creates a provider for the canonical atom-aligned pLDDT confidence column.
226    ///
227    /// # Errors
228    ///
229    /// Rejects structures without lowered pLDDT, row misalignment, or identity overflow.
230    pub fn plddt(
231        dataset: DatasetId,
232        first_chunk: ChunkId,
233        structure: Structure,
234    ) -> Result<Self, ProviderError> {
235        Self::new(
236            dataset,
237            first_chunk,
238            structure,
239            crate::annotation::PLDDT_ANNOTATION,
240        )
241    }
242
243    /// Creates a property provider over the structure's existing chunk layout.
244    ///
245    /// # Errors
246    ///
247    /// Rejects absent properties, row misalignment, or identity overflow.
248    pub fn new(
249        dataset: DatasetId,
250        first_chunk: ChunkId,
251        structure: Structure,
252        name: impl Into<Arc<str>>,
253    ) -> Result<Self, ProviderError> {
254        let name = name.into();
255        let property = structure.annotations().get(&name).ok_or_else(|| {
256            ProviderError::PropertyUnavailable {
257                name: Box::from(name.as_ref()),
258            }
259        })?;
260        if property.len() != structure.atom_count() {
261            return Err(ProviderError::PropertyLengthMismatch {
262                name: Box::from(name.as_ref()),
263                expected: structure.atom_count(),
264                actual: property.len(),
265            });
266        }
267        let chunk_count = u64::try_from(structure.data().chunks.len()).map_err(|_| {
268            ProviderError::IdentityOverflow {
269                dataset,
270                field: "property chunk count",
271            }
272        })?;
273        let descriptor = DatasetDescriptor::source_defined(
274            dataset,
275            PayloadKind::Property,
276            u64::from(structure.atom_count()),
277            chunk_count,
278            first_chunk,
279            TARGET_CHUNK_ATOMS,
280        )?;
281        Ok(Self {
282            structure,
283            name,
284            dataset: descriptor,
285        })
286    }
287
288    /// Compact dataset metadata.
289    #[must_use]
290    pub const fn dataset(&self) -> DatasetDescriptor {
291        self.dataset
292    }
293
294    /// Produces one property payload without copying its values.
295    ///
296    /// # Errors
297    ///
298    /// Rejects chunk identities outside this provider.
299    pub fn chunk(&self, chunk: ChunkId) -> Result<PropertyChunk, ProviderError> {
300        let (index, descriptor) = storage_descriptor(&self.dataset, &self.structure, chunk)?;
301        PropertyChunk::shared(
302            descriptor,
303            self.structure.clone(),
304            Arc::clone(&self.name),
305            index,
306        )
307    }
308}
309
310/// Zero-copy provider for one coordinate frame over an existing chunk layout.
311#[derive(Clone, Debug)]
312pub struct FrameChunkProvider {
313    structure: Structure,
314    model: crate::ModelIndex,
315    storage_chunks: std::ops::Range<usize>,
316    dataset: DatasetDescriptor,
317}
318
319impl FrameChunkProvider {
320    /// Creates a frame provider for one directly stored model.
321    ///
322    /// # Errors
323    ///
324    /// Rejects ragged or absent model storage and identity overflow.
325    pub fn new(
326        dataset: DatasetId,
327        first_chunk: ChunkId,
328        structure: Structure,
329        model: crate::ModelIndex,
330    ) -> Result<Self, ProviderError> {
331        let Some(coordinates) = structure.data().coords.block(model) else {
332            return Err(ProviderError::ModelUnavailable { model: model.get() });
333        };
334        let storage_chunks = model_chunk_range(&structure, model);
335        let chunk_count =
336            u64::try_from(storage_chunks.len()).map_err(|_| ProviderError::IdentityOverflow {
337                dataset,
338                field: "frame chunk count",
339            })?;
340        let descriptor = DatasetDescriptor::source_defined(
341            dataset,
342            PayloadKind::Frame,
343            u64::from(coordinates.len()),
344            chunk_count,
345            first_chunk,
346            TARGET_CHUNK_ATOMS,
347        )?;
348        Ok(Self {
349            structure,
350            model,
351            storage_chunks,
352            dataset: descriptor,
353        })
354    }
355
356    /// Compact dataset metadata.
357    #[must_use]
358    pub const fn dataset(&self) -> DatasetDescriptor {
359        self.dataset
360    }
361
362    /// Produces one coordinate payload sharing the frame's aligned storage.
363    ///
364    /// # Errors
365    ///
366    /// Rejects chunk identities outside this provider or inconsistent models.
367    pub fn chunk(&self, chunk: ChunkId) -> Result<FrameChunk, ProviderError> {
368        let ordinal = self.dataset.chunk_ordinal(chunk)?;
369        let local = usize::try_from(ordinal).map_err(|_| ProviderError::IdentityOverflow {
370            dataset: self.dataset.id(),
371            field: "storage chunk index",
372        })?;
373        let index = self.storage_chunks.start.checked_add(local).ok_or(
374            ProviderError::IdentityOverflow {
375                dataset: self.dataset.id(),
376                field: "storage chunk index",
377            },
378        )?;
379        if index >= self.storage_chunks.end {
380            return Err(ProviderError::UnknownChunk {
381                dataset: self.dataset.id(),
382                chunk,
383            });
384        }
385        let storage = &self.structure.data().chunks[index];
386        let range = storage.atoms();
387        let descriptor = ChunkDescriptor::new(
388            self.dataset.id(),
389            chunk,
390            super::LogicalRow::new(u64::from(range.start)),
391            storage.len(),
392        )?;
393        let coordinates = self
394            .structure
395            .data()
396            .coords
397            .block(self.model)
398            .cloned()
399            .ok_or(ProviderError::ModelUnavailable {
400                model: self.model.get(),
401            })?;
402        FrameChunk::shared(descriptor, coordinates, range)
403    }
404}
405
406fn model_chunk_range(structure: &Structure, model: crate::ModelIndex) -> std::ops::Range<usize> {
407    let chunks = structure.data().chunks.as_slice();
408    let start = chunks.partition_point(|chunk| chunk.model() < model.get());
409    let end = chunks.partition_point(|chunk| chunk.model() <= model.get());
410    start..end
411}
412
413#[cfg(test)]
414#[path = "source_tests.rs"]
415mod tests;
416
417#[cfg(test)]
418#[path = "bond_source_tests.rs"]
419mod bond_tests;
molframe_core/provider/source.rs

molframe_core/provider/
source.rs
