Skip to main content

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 validate_atom_endpoint_range(
120    dataset: DatasetId,
121    first: LogicalRow,
122    rows: u32,
123) -> Result<(), ProviderError> {
124    if rows == 0 {
125        return Ok(());
126    }
127    first
128        .get()
129        .checked_add(u64::from(rows - 1))
130        .map(|_| ())
131        .ok_or(ProviderError::IdentityOverflow {
132            dataset,
133            field: "atom endpoint range",
134        })
135}
136
137/// Zero-copy provider for structure chunks already present in one snapshot.
138#[derive(Clone, Debug)]
139pub struct StructureChunkProvider {
140    structure: Structure,
141    dataset: DatasetDescriptor,
142}
143
144impl StructureChunkProvider {
145    /// Creates a provider whose chunk namespace begins at `first_chunk`.
146    ///
147    /// # Errors
148    ///
149    /// Returns a typed error when chunk count or identity arithmetic overflows.
150    pub fn new(
151        dataset: DatasetId,
152        first_chunk: ChunkId,
153        structure: Structure,
154    ) -> Result<Self, ProviderError> {
155        let chunk_count = u64::try_from(structure.data().chunks.len()).map_err(|_| {
156            ProviderError::IdentityOverflow {
157                dataset,
158                field: "structure chunk count",
159            }
160        })?;
161        let descriptor = DatasetDescriptor::source_defined(
162            dataset,
163            PayloadKind::Structure,
164            u64::from(structure.atom_count()),
165            chunk_count,
166            first_chunk,
167            TARGET_CHUNK_ATOMS,
168        )?;
169        Ok(Self {
170            structure,
171            dataset: descriptor,
172        })
173    }
174
175    /// Compact dataset metadata.
176    #[must_use]
177    pub const fn dataset(&self) -> DatasetDescriptor {
178        self.dataset
179    }
180
181    /// Produces one shared structure payload in constant time.
182    ///
183    /// # Errors
184    ///
185    /// Rejects chunk identities outside this provider.
186    pub fn chunk(&self, chunk: ChunkId) -> Result<StructureChunk, ProviderError> {
187        let ordinal = self.dataset.chunk_ordinal(chunk)?;
188        let index = usize::try_from(ordinal).map_err(|_| ProviderError::IdentityOverflow {
189            dataset: self.dataset.id(),
190            field: "storage chunk index",
191        })?;
192        let storage = self
193            .structure
194            .data()
195            .chunks
196            .get(index)
197            .ok_or(ProviderError::StorageChunkUnavailable { index })?;
198        let descriptor = ChunkDescriptor::new(
199            self.dataset.id(),
200            chunk,
201            super::LogicalRow::new(u64::from(storage.atoms().start)),
202            storage.len(),
203        )?;
204        StructureChunk::shared(descriptor, self.structure.clone(), index)
205    }
206}
207
208/// Zero-copy provider for one typed atom property.
209#[derive(Clone, Debug)]
210pub struct PropertyChunkProvider {
211    structure: Structure,
212    name: Arc<str>,
213    dataset: DatasetDescriptor,
214}
215
216impl PropertyChunkProvider {
217    /// Creates a provider for the canonical atom-aligned pLDDT confidence column.
218    ///
219    /// # Errors
220    ///
221    /// Rejects structures without lowered pLDDT, row misalignment, or identity overflow.
222    pub fn plddt(
223        dataset: DatasetId,
224        first_chunk: ChunkId,
225        structure: Structure,
226    ) -> Result<Self, ProviderError> {
227        Self::new(
228            dataset,
229            first_chunk,
230            structure,
231            crate::annotation::PLDDT_ANNOTATION,
232        )
233    }
234
235    /// Creates a property provider over the structure's existing chunk layout.
236    ///
237    /// # Errors
238    ///
239    /// Rejects absent properties, row misalignment, or identity overflow.
240    pub fn new(
241        dataset: DatasetId,
242        first_chunk: ChunkId,
243        structure: Structure,
244        name: impl Into<Arc<str>>,
245    ) -> Result<Self, ProviderError> {
246        let name = name.into();
247        let property = structure.annotations().get(&name).ok_or_else(|| {
248            ProviderError::PropertyUnavailable {
249                name: Box::from(name.as_ref()),
250            }
251        })?;
252        if property.len() != structure.atom_count() {
253            return Err(ProviderError::PropertyLengthMismatch {
254                name: Box::from(name.as_ref()),
255                expected: structure.atom_count(),
256                actual: property.len(),
257            });
258        }
259        let chunk_count = u64::try_from(structure.data().chunks.len()).map_err(|_| {
260            ProviderError::IdentityOverflow {
261                dataset,
262                field: "property chunk count",
263            }
264        })?;
265        let descriptor = DatasetDescriptor::source_defined(
266            dataset,
267            PayloadKind::Property,
268            u64::from(structure.atom_count()),
269            chunk_count,
270            first_chunk,
271            TARGET_CHUNK_ATOMS,
272        )?;
273        Ok(Self {
274            structure,
275            name,
276            dataset: descriptor,
277        })
278    }
279
280    /// Compact dataset metadata.
281    #[must_use]
282    pub const fn dataset(&self) -> DatasetDescriptor {
283        self.dataset
284    }
285
286    /// Produces one property payload without copying its values.
287    ///
288    /// # Errors
289    ///
290    /// Rejects chunk identities outside this provider.
291    pub fn chunk(&self, chunk: ChunkId) -> Result<PropertyChunk, ProviderError> {
292        let ordinal = self.dataset.chunk_ordinal(chunk)?;
293        let index = usize::try_from(ordinal).map_err(|_| ProviderError::IdentityOverflow {
294            dataset: self.dataset.id(),
295            field: "storage chunk index",
296        })?;
297        let storage = self
298            .structure
299            .data()
300            .chunks
301            .get(index)
302            .ok_or(ProviderError::StorageChunkUnavailable { index })?;
303        let descriptor = ChunkDescriptor::new(
304            self.dataset.id(),
305            chunk,
306            super::LogicalRow::new(u64::from(storage.atoms().start)),
307            storage.len(),
308        )?;
309        PropertyChunk::shared(
310            descriptor,
311            self.structure.clone(),
312            Arc::clone(&self.name),
313            index,
314        )
315    }
316}
317
318/// Zero-copy provider for one coordinate frame over an existing chunk layout.
319#[derive(Clone, Debug)]
320pub struct FrameChunkProvider {
321    structure: Structure,
322    model: crate::ModelIndex,
323    storage_chunks: std::ops::Range<usize>,
324    dataset: DatasetDescriptor,
325}
326
327impl FrameChunkProvider {
328    /// Creates a frame provider for one directly stored model.
329    ///
330    /// # Errors
331    ///
332    /// Rejects ragged or absent model storage and identity overflow.
333    pub fn new(
334        dataset: DatasetId,
335        first_chunk: ChunkId,
336        structure: Structure,
337        model: crate::ModelIndex,
338    ) -> Result<Self, ProviderError> {
339        let Some(coordinates) = structure.data().coords.block(model) else {
340            return Err(ProviderError::ModelUnavailable { model: model.get() });
341        };
342        let storage_chunks = model_chunk_range(&structure, model);
343        let chunk_count =
344            u64::try_from(storage_chunks.len()).map_err(|_| ProviderError::IdentityOverflow {
345                dataset,
346                field: "frame chunk count",
347            })?;
348        let descriptor = DatasetDescriptor::source_defined(
349            dataset,
350            PayloadKind::Frame,
351            u64::from(coordinates.len()),
352            chunk_count,
353            first_chunk,
354            TARGET_CHUNK_ATOMS,
355        )?;
356        Ok(Self {
357            structure,
358            model,
359            storage_chunks,
360            dataset: descriptor,
361        })
362    }
363
364    /// Compact dataset metadata.
365    #[must_use]
366    pub const fn dataset(&self) -> DatasetDescriptor {
367        self.dataset
368    }
369
370    /// Produces one coordinate payload sharing the frame's aligned storage.
371    ///
372    /// # Errors
373    ///
374    /// Rejects chunk identities outside this provider or inconsistent models.
375    pub fn chunk(&self, chunk: ChunkId) -> Result<FrameChunk, ProviderError> {
376        let ordinal = self.dataset.chunk_ordinal(chunk)?;
377        let local = usize::try_from(ordinal).map_err(|_| ProviderError::IdentityOverflow {
378            dataset: self.dataset.id(),
379            field: "storage chunk index",
380        })?;
381        let index = self.storage_chunks.start.checked_add(local).ok_or(
382            ProviderError::IdentityOverflow {
383                dataset: self.dataset.id(),
384                field: "storage chunk index",
385            },
386        )?;
387        if index >= self.storage_chunks.end {
388            return Err(ProviderError::UnknownChunk {
389                dataset: self.dataset.id(),
390                chunk,
391            });
392        }
393        let storage = &self.structure.data().chunks[index];
394        let range = storage.atoms();
395        let descriptor = ChunkDescriptor::new(
396            self.dataset.id(),
397            chunk,
398            super::LogicalRow::new(u64::from(range.start)),
399            storage.len(),
400        )?;
401        let coordinates = self
402            .structure
403            .data()
404            .coords
405            .block(self.model)
406            .cloned()
407            .ok_or(ProviderError::ModelUnavailable {
408                model: self.model.get(),
409            })?;
410        FrameChunk::shared(descriptor, coordinates, range)
411    }
412}
413
414fn model_chunk_range(structure: &Structure, model: crate::ModelIndex) -> std::ops::Range<usize> {
415    let chunks = structure.data().chunks.as_slice();
416    let start = chunks.partition_point(|chunk| chunk.model() < model.get());
417    let end = chunks.partition_point(|chunk| chunk.model() <= model.get());
418    start..end
419}
420
421#[cfg(test)]
422#[path = "source_tests.rs"]
423mod tests;
424
425#[cfg(test)]
426#[path = "bond_source_tests.rs"]
427mod bond_tests;