1use crate::column::Presence;
4use crate::execution::Batch;
5use crate::optional::OptionalSymbol;
6use crate::provider::{ChunkDescriptor, ProviderError};
7use crate::symbol::{AltId, DictionaryFull, Interner, SymbolId};
8use crate::{Diagnostic, Element};
9use std::fmt;
10
11#[derive(Debug)]
13pub enum StructureBatchError {
14 Diagnostic(Diagnostic),
16 DemandTooSmall {
18 required: usize,
20 available: usize,
22 },
23 RecordExceedsBudget {
25 required: usize,
27 available: usize,
29 },
30 Identity(ProviderError),
32 Memory(crate::MemoryBudgetError),
34 DictionaryFull,
36}
37
38impl fmt::Display for StructureBatchError {
39 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
40 match self {
41 Self::Diagnostic(error) => error.fmt(formatter),
42 Self::DemandTooSmall {
43 required,
44 available,
45 } => write!(
46 formatter,
47 "record requires {required} bytes but demand permits {available}"
48 ),
49 Self::RecordExceedsBudget {
50 required,
51 available,
52 } => write!(
53 formatter,
54 "record requires {required} bytes but budget permits {available}"
55 ),
56 Self::Identity(error) => error.fmt(formatter),
57 Self::Memory(error) => error.fmt(formatter),
58 Self::DictionaryFull => formatter.write_str("batch identifier dictionary is full"),
59 }
60 }
61}
62
63impl std::error::Error for StructureBatchError {}
64
65impl From<Diagnostic> for StructureBatchError {
66 fn from(error: Diagnostic) -> Self {
67 Self::Diagnostic(error)
68 }
69}
70
71impl From<ProviderError> for StructureBatchError {
72 fn from(error: ProviderError) -> Self {
73 Self::Identity(error)
74 }
75}
76
77impl From<DictionaryFull> for StructureBatchError {
78 fn from(_: DictionaryFull) -> Self {
79 Self::DictionaryFull
80 }
81}
82
83impl From<crate::MemoryBudgetError> for StructureBatchError {
84 fn from(error: crate::MemoryBudgetError) -> Self {
85 Self::Memory(error)
86 }
87}
88
89#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
91pub enum ContinuityLevel {
92 #[default]
94 None,
95 Model,
97 Chain,
99 Residue,
101}
102
103#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
105pub struct BatchContinuity {
106 pub before: ContinuityLevel,
108 pub after: ContinuityLevel,
110}
111
112#[derive(Clone, Copy, Debug)]
114pub struct StructureAtomRecord<'a> {
115 pub model: i32,
117 pub chain: &'a str,
119 pub component: &'a str,
121 pub sequence: i32,
123 pub insertion: &'a str,
125 pub atom: &'a str,
127 pub alternate: &'a str,
129 pub element: Element,
131 pub position: Option<[f32; 3]>,
133 pub occupancy: (f32, Presence),
135 pub b_factor: (f32, Presence),
137 pub formal_charge: (i8, Presence),
139 pub atom_site_id: u32,
141 pub heterogen: bool,
143}
144
145#[derive(Clone, Debug)]
147pub struct StructureBatch {
148 pub(super) descriptor: ChunkDescriptor,
149 pub(super) continuity: BatchContinuity,
150 pub(super) dictionary: Interner,
151 pub(super) model: Vec<i32>,
152 pub(super) chain: Vec<SymbolId>,
153 pub(super) component: Vec<SymbolId>,
154 pub(super) sequence: Vec<i32>,
155 pub(super) insertion: Vec<OptionalSymbol>,
156 pub(super) atom: Vec<SymbolId>,
157 pub(super) alternate: Vec<AltId>,
158 pub(super) element: Vec<u8>,
159 pub(super) position: Vec<[f32; 3]>,
160 pub(super) coordinate_presence: Vec<Presence>,
161 pub(super) occupancy: Vec<f32>,
162 pub(super) occupancy_presence: Vec<Presence>,
163 pub(super) b_factor: Vec<f32>,
164 pub(super) b_factor_presence: Vec<Presence>,
165 pub(super) formal_charge: Vec<i8>,
166 pub(super) formal_charge_presence: Vec<Presence>,
167 pub(super) atom_site_id: Vec<u32>,
168 pub(super) heterogen: Vec<u8>,
169 pub(super) diagnostics: Vec<Diagnostic>,
170}
171
172impl StructureBatch {
173 #[must_use]
176 pub fn can_reuse(&self, rows: u32, max_retained_bytes: usize) -> bool {
177 let rows = rows as usize;
178 self.model.capacity() >= rows && self.retained_bytes() <= max_retained_bytes
179 }
180
181 #[must_use]
183 pub const fn descriptor(&self) -> ChunkDescriptor {
184 self.descriptor
185 }
186
187 #[must_use]
189 pub const fn continuity(&self) -> BatchContinuity {
190 self.continuity
191 }
192
193 #[must_use]
195 pub const fn dictionary(&self) -> &Interner {
196 &self.dictionary
197 }
198
199 #[must_use]
201 pub fn models(&self) -> &[i32] {
202 &self.model
203 }
204
205 #[must_use]
207 pub fn chains(&self) -> &[SymbolId] {
208 &self.chain
209 }
210
211 #[must_use]
213 pub fn components(&self) -> &[SymbolId] {
214 &self.component
215 }
216
217 #[must_use]
219 pub fn sequences(&self) -> &[i32] {
220 &self.sequence
221 }
222
223 #[must_use]
225 pub fn insertions(&self) -> &[OptionalSymbol] {
226 &self.insertion
227 }
228
229 #[must_use]
231 pub fn atoms(&self) -> &[SymbolId] {
232 &self.atom
233 }
234
235 #[must_use]
237 pub fn alternates(&self) -> &[AltId] {
238 &self.alternate
239 }
240
241 #[must_use]
243 pub fn elements(&self) -> &[u8] {
244 &self.element
245 }
246
247 #[must_use]
249 pub fn positions(&self) -> &[[f32; 3]] {
250 &self.position
251 }
252
253 #[must_use]
255 pub fn coordinate_presence(&self) -> &[Presence] {
256 &self.coordinate_presence
257 }
258
259 #[must_use]
261 pub fn occupancies(&self) -> (&[f32], &[Presence]) {
262 (&self.occupancy, &self.occupancy_presence)
263 }
264
265 #[must_use]
267 pub fn b_factors(&self) -> (&[f32], &[Presence]) {
268 (&self.b_factor, &self.b_factor_presence)
269 }
270
271 #[must_use]
273 pub fn formal_charges(&self) -> (&[i8], &[Presence]) {
274 (&self.formal_charge, &self.formal_charge_presence)
275 }
276
277 #[must_use]
279 pub fn atom_site_ids(&self) -> &[u32] {
280 &self.atom_site_id
281 }
282
283 #[must_use]
285 pub fn heterogens(&self) -> &[u8] {
286 &self.heterogen
287 }
288
289 #[must_use]
291 pub fn diagnostics(&self) -> &[Diagnostic] {
292 &self.diagnostics
293 }
294}
295
296impl Batch for StructureBatch {
297 fn rows(&self) -> usize {
298 self.model.len()
299 }
300
301 fn retained_bytes(&self) -> usize {
302 let mut bytes = self.dictionary.retained_bytes();
303 bytes = bytes.saturating_add(self.model.capacity_bytes());
304 bytes = bytes.saturating_add(self.chain.capacity_bytes());
305 bytes = bytes.saturating_add(self.component.capacity_bytes());
306 bytes = bytes.saturating_add(self.sequence.capacity_bytes());
307 bytes = bytes.saturating_add(self.insertion.capacity_bytes());
308 bytes = bytes.saturating_add(self.atom.capacity_bytes());
309 bytes = bytes.saturating_add(self.alternate.capacity_bytes());
310 bytes = bytes.saturating_add(self.element.capacity_bytes());
311 bytes = bytes.saturating_add(self.position.capacity_bytes());
312 bytes = bytes.saturating_add(self.coordinate_presence.capacity_bytes());
313 bytes = bytes.saturating_add(self.occupancy.capacity_bytes());
314 bytes = bytes.saturating_add(self.occupancy_presence.capacity_bytes());
315 bytes = bytes.saturating_add(self.b_factor.capacity_bytes());
316 bytes = bytes.saturating_add(self.b_factor_presence.capacity_bytes());
317 bytes = bytes.saturating_add(self.formal_charge.capacity_bytes());
318 bytes = bytes.saturating_add(self.formal_charge_presence.capacity_bytes());
319 bytes = bytes.saturating_add(self.atom_site_id.capacity_bytes());
320 bytes = bytes.saturating_add(self.heterogen.capacity_bytes());
321 bytes = bytes.saturating_add(self.diagnostics.capacity_bytes());
322 bytes.saturating_add(
323 self.diagnostics
324 .iter()
325 .map(Diagnostic::retained_bytes)
326 .sum::<usize>(),
327 )
328 }
329}
330
331trait VecCapacityBytes {
332 fn capacity_bytes(&self) -> usize;
333}
334
335impl<T> VecCapacityBytes for Vec<T> {
336 fn capacity_bytes(&self) -> usize {
337 self.capacity().saturating_mul(std::mem::size_of::<T>())
338 }
339}