1use 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
11pub const TARGET_CHUNK_BONDS: u32 = 65_536;
13
14#[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 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 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 #[must_use]
86 pub const fn dataset(&self) -> DatasetDescriptor {
87 self.dataset
88 }
89
90 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#[derive(Clone, Debug)]
139pub struct StructureChunkProvider {
140 structure: Structure,
141 dataset: DatasetDescriptor,
142}
143
144impl StructureChunkProvider {
145 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 #[must_use]
177 pub const fn dataset(&self) -> DatasetDescriptor {
178 self.dataset
179 }
180
181 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#[derive(Clone, Debug)]
210pub struct PropertyChunkProvider {
211 structure: Structure,
212 name: Arc<str>,
213 dataset: DatasetDescriptor,
214}
215
216impl PropertyChunkProvider {
217 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 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 #[must_use]
282 pub const fn dataset(&self) -> DatasetDescriptor {
283 self.dataset
284 }
285
286 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#[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 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 #[must_use]
366 pub const fn dataset(&self) -> DatasetDescriptor {
367 self.dataset
368 }
369
370 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;