document_svg/document/
pdb.rs

1//! Bounded Protein Data Bank coordinate previews.
2//!
3//! This reader consumes fixed-column ATOM/HETATM coordinates and optional
4//! CONECT records from legacy PDB text files. Each MODEL becomes an SVG page;
5//! coordinates are projected onto XY and chemical interpretation is deliberately
6//! limited to the element labels supplied by the file.
7
8use std::collections::{HashMap, HashSet};
9use std::path::Path;
10
11use crate::convert::{ConvertOptions, PageConsumer, read_limited_file};
12use crate::document::molfile::{Atom, Bond, Molecule, render_molecule};
13use crate::error::{Error, Result};
14
15const MAX_PDB_INPUT_BYTES: u64 = 64 * 1024 * 1024;
16const MAX_PDB_LINES: usize = 1_000_000;
17const MAX_PDB_LINE_BYTES: usize = 1024 * 1024;
18const MAX_PDB_MODELS: usize = 1_000;
19const MAX_PDB_ATOMS_PER_MODEL: usize = 100_000;
20const MAX_PDB_TOTAL_ATOMS: usize = 500_000;
21const MAX_PDB_BONDS_PER_MODEL: usize = 200_000;
22const MAX_PDB_COORDINATE: f64 = 1_000_000.0;
23
24struct ModelBuilder {
25    number: usize,
26    title: String,
27    atoms: Vec<Atom>,
28    serials: HashMap<String, usize>,
29    conect: Vec<(String, String)>,
30    warnings: Vec<String>,
31    has_nonzero_z: bool,
32    skipped_altlocs: usize,
33}
34
35impl ModelBuilder {
36    fn new(number: usize, title: String) -> Self {
37        Self {
38            number,
39            title,
40            atoms: Vec::new(),
41            serials: HashMap::new(),
42            conect: Vec::new(),
43            warnings: Vec::new(),
44            has_nonzero_z: false,
45            skipped_altlocs: 0,
46        }
47    }
48
49    fn finish(self) -> Result<Molecule> {
50        let mut bonds = Vec::new();
51        let mut seen = HashSet::new();
52        for (from_serial, to_serial) in self.conect {
53            let Some(&from) = self.serials.get(&from_serial) else {
54                continue;
55            };
56            let Some(&to) = self.serials.get(&to_serial) else {
57                continue;
58            };
59            if from == to {
60                continue;
61            }
62            let key = (from.min(to), from.max(to));
63            if seen.insert(key) {
64                if bonds.len() >= MAX_PDB_BONDS_PER_MODEL {
65                    return Err(Error::LimitExceeded(format!(
66                        "PDB model exceeds {MAX_PDB_BONDS_PER_MODEL} CONECT bonds"
67                    )));
68                }
69                bonds.push(Bond {
70                    from,
71                    to,
72                    order: 1,
73                    stereo: 0,
74                });
75            }
76        }
77        let mut warnings = self.warnings;
78        if self.has_nonzero_z {
79            warnings.push("PDB 3D coordinates are projected onto the XY plane".into());
80        }
81        if self.skipped_altlocs > 0 {
82            warnings.push(format!(
83                "PDB alternate atom locations other than blank/A were omitted ({})",
84                self.skipped_altlocs
85            ));
86        }
87        if !self.atoms.is_empty() && bonds.is_empty() {
88            warnings.push("PDB CONECT records were absent; atoms are shown without bonds".into());
89        }
90        Ok(Molecule {
91            title: if self.title.is_empty() {
92                format!("PDB model {}", self.number)
93            } else if self.number > 1 {
94                format!("{} (model {})", self.title, self.number)
95            } else {
96                self.title
97            },
98            atoms: self.atoms,
99            bonds,
100            warnings,
101        })
102    }
103}
104
105pub(crate) fn looks_like_prefix(bytes: &[u8]) -> bool {
106    let text = String::from_utf8_lossy(bytes);
107    text.lines().take(32).any(|line| {
108        let record = line.get(..6).unwrap_or(line).trim();
109        matches!(record, "ATOM" | "HETATM" | "MODEL" | "HEADER")
110    })
111}
112
113pub(crate) fn convert(
114    path: &Path,
115    options: &ConvertOptions,
116    sink: &mut dyn PageConsumer,
117) -> Result<Vec<String>> {
118    let bytes = read_limited_file(
119        path,
120        options.max_input_bytes.min(MAX_PDB_INPUT_BYTES),
121        "PDB input",
122    )?;
123    let text = std::str::from_utf8(&bytes).map_err(|error| {
124        Error::InvalidInput(format!("PDB input is not UTF-8 or ASCII: {error}"))
125    })?;
126    let models = parse_models(text)?;
127    if models.is_empty() {
128        return Err(Error::InvalidInput(
129            "PDB input contains no ATOM or HETATM records".into(),
130        ));
131    }
132    if options.max_pages == 0 {
133        return Err(Error::LimitExceeded(
134            "PDB conversion requires at least one page; max_pages is zero".into(),
135        ));
136    }
137    if models.len() > options.max_pages {
138        return Err(Error::LimitExceeded(format!(
139            "PDB contains {} models but max_pages is {}",
140            models.len(),
141            options.max_pages
142        )));
143    }
144    let mut warnings = Vec::new();
145    for (index, model) in models.iter().enumerate() {
146        let page = render_molecule(model, index + 1, "pdb")?;
147        warnings.extend(page.warnings.iter().cloned());
148        sink.consume(page)?;
149    }
150    warnings.sort();
151    warnings.dedup();
152    Ok(warnings)
153}
154
155fn parse_models(text: &str) -> Result<Vec<Molecule>> {
156    let mut models = Vec::new();
157    let mut current = None::<ModelBuilder>;
158    let mut next_model = 1usize;
159    let mut total_atoms = 0usize;
160    let mut line_count = 0usize;
161    let mut global_title = String::new();
162    for line in text.lines() {
163        line_count = line_count.saturating_add(1);
164        if line_count > MAX_PDB_LINES {
165            return Err(Error::LimitExceeded(format!(
166                "PDB input exceeds {MAX_PDB_LINES} lines"
167            )));
168        }
169        if line.len() > MAX_PDB_LINE_BYTES {
170            return Err(Error::LimitExceeded(format!(
171                "PDB line exceeds {MAX_PDB_LINE_BYTES} bytes"
172            )));
173        }
174        let record = line.get(..6).unwrap_or(line).trim();
175        match record {
176            "HEADER" => {
177                if global_title.is_empty() {
178                    global_title = clean_field(
179                        line.get(10..50)
180                            .or_else(|| line.get(10..))
181                            .unwrap_or_default(),
182                    );
183                }
184            }
185            "TITLE" => {
186                let part = clean_field(
187                    line.get(10..80)
188                        .or_else(|| line.get(10..))
189                        .unwrap_or_default(),
190                );
191                if !part.is_empty() {
192                    if !global_title.is_empty() {
193                        global_title.push(' ');
194                    }
195                    global_title.push_str(&part);
196                }
197            }
198            "MODEL" => {
199                if let Some(builder) = current.take() {
200                    models.push(builder.finish()?);
201                }
202                if models.len() >= MAX_PDB_MODELS {
203                    return Err(Error::LimitExceeded(format!(
204                        "PDB exceeds {MAX_PDB_MODELS} models"
205                    )));
206                }
207                let number = parse_optional_usize(line.get(10..14).unwrap_or_default())
208                    .unwrap_or(next_model);
209                next_model = number.saturating_add(1);
210                current = Some(ModelBuilder::new(number, global_title.clone()));
211            }
212            "ENDMDL" => {
213                if let Some(builder) = current.take() {
214                    models.push(builder.finish()?);
215                }
216            }
217            "ATOM" | "HETATM" => {
218                if current.is_none() {
219                    current = Some(ModelBuilder::new(next_model, global_title.clone()));
220                    next_model = next_model.saturating_add(1);
221                }
222                let builder = current.as_mut().expect("PDB model initialized");
223                if builder.atoms.len() >= MAX_PDB_ATOMS_PER_MODEL {
224                    return Err(Error::LimitExceeded(format!(
225                        "PDB model exceeds {MAX_PDB_ATOMS_PER_MODEL} atoms"
226                    )));
227                }
228                let alt = line.as_bytes().get(16).copied().unwrap_or(b' ');
229                if alt != b' ' && alt != b'A' {
230                    builder.skipped_altlocs = builder.skipped_altlocs.saturating_add(1);
231                    continue;
232                }
233                let serial = fixed_field(line, 6, 11, "atom serial")?;
234                let x = parse_coordinate(fixed_field(line, 30, 38, "x coordinate")?, "x")?;
235                let y = parse_coordinate(fixed_field(line, 38, 46, "y coordinate")?, "y")?;
236                let z = parse_coordinate(fixed_field(line, 46, 54, "z coordinate")?, "z")?;
237                let atom_name = line.get(12..16).unwrap_or_default();
238                let element = element_name(line.get(76..78).unwrap_or_default(), atom_name);
239                let index = builder.atoms.len();
240                builder.serials.entry(serial).or_insert(index);
241                builder.has_nonzero_z |= z.abs() > 1e-6;
242                builder.atoms.push(Atom {
243                    x,
244                    y,
245                    element,
246                    charge: 0,
247                    isotope: None,
248                });
249                total_atoms = total_atoms.saturating_add(1);
250                if total_atoms > MAX_PDB_TOTAL_ATOMS {
251                    return Err(Error::LimitExceeded(format!(
252                        "PDB exceeds {MAX_PDB_TOTAL_ATOMS} total atoms"
253                    )));
254                }
255            }
256            "CONECT" => {
257                if let Some(builder) = current.as_mut() {
258                    let source = fixed_field(line, 6, 11, "CONECT source")?;
259                    for start in [11usize, 16, 21, 26, 31] {
260                        if start >= line.len() {
261                            break;
262                        }
263                        let target = line
264                            .get(start..start.saturating_add(5))
265                            .unwrap_or_default()
266                            .trim();
267                        if !target.is_empty() {
268                            builder.conect.push((source.clone(), target.to_owned()));
269                        }
270                    }
271                }
272            }
273            "END" => break,
274            _ => {}
275        }
276    }
277    if let Some(builder) = current.take() {
278        models.push(builder.finish()?);
279    }
280    if models.len() > MAX_PDB_MODELS {
281        return Err(Error::LimitExceeded(format!(
282            "PDB exceeds {MAX_PDB_MODELS} models"
283        )));
284    }
285    Ok(models)
286}
287
288fn fixed_field(line: &str, start: usize, end: usize, label: &str) -> Result<String> {
289    let value = line
290        .get(start..end.min(line.len()))
291        .ok_or_else(|| Error::InvalidInput(format!("PDB {label} field is missing")))?
292        .trim();
293    if value.is_empty() {
294        return Err(Error::InvalidInput(format!("PDB {label} field is empty")));
295    }
296    Ok(value.to_owned())
297}
298
299fn parse_coordinate(value: String, label: &str) -> Result<f64> {
300    let parsed = value
301        .parse::<f64>()
302        .map_err(|_| Error::InvalidInput(format!("PDB {label} coordinate is invalid")))?;
303    if !parsed.is_finite() || parsed.abs() > MAX_PDB_COORDINATE {
304        return Err(Error::InvalidInput(format!(
305            "PDB {label} coordinate exceeds ±{MAX_PDB_COORDINATE}"
306        )));
307    }
308    Ok(parsed)
309}
310
311fn element_name(field: &str, atom_name: &str) -> String {
312    let candidate = if field.trim().is_empty() {
313        atom_name
314            .chars()
315            .filter(|character| character.is_ascii_alphabetic())
316            .take(2)
317            .collect::<String>()
318    } else {
319        field
320            .trim()
321            .chars()
322            .filter(|c| c.is_ascii_alphabetic())
323            .take(2)
324            .collect()
325    };
326    let mut chars = candidate.chars();
327    let Some(first) = chars.next() else {
328        return "X".into();
329    };
330    let mut output = first.to_ascii_uppercase().to_string();
331    if let Some(second) = chars.next() {
332        output.push(second.to_ascii_lowercase());
333    }
334    output
335}
336
337fn clean_field(value: &str) -> String {
338    value
339        .trim()
340        .chars()
341        .filter(|character| !character.is_control())
342        .take(240)
343        .collect()
344}
345
346fn parse_optional_usize(value: &str) -> Option<usize> {
347    value.trim().parse::<usize>().ok()
348}
349
350#[cfg(test)]
351mod tests {
352    use super::{looks_like_prefix, parse_models};
353
354    #[test]
355    fn recognizes_pdb_records_without_broad_text_sniffing() {
356        assert!(looks_like_prefix(
357            b"HEADER    SAMPLE\nATOM      1  CA  ALA A   1       0.000   0.000   0.000\n"
358        ));
359        assert!(!looks_like_prefix(b"This prose mentions ATOM as a word.\n"));
360    }
361
362    #[test]
363    fn parses_fixed_columns_and_model_title() {
364        let text = "HEADER    SAMPLE PDB PREVIEW\nTITLE     Two model test\nMODEL        1\nATOM      1  CA  ALA A   1       0.000   0.000   0.000  1.00 10.00           C  \nENDMDL\n";
365        let models = parse_models(text).unwrap();
366        assert_eq!(models.len(), 1);
367        assert!(
368            models[0].title.contains("SAMPLE PDB PREVIEW"),
369            "title={:?}",
370            models[0].title
371        );
372        assert_eq!(models[0].atoms.len(), 1);
373    }
374}
document_svg/document/pdb.rs

document_svg/document/
pdb.rs
