1use super::super::input::InputBuffer;
4use crate::diagnostic::{Code, Diagnostic};
5
6const PDBML_PROBE_WINDOWS: usize = 2_048;
7const MESSAGEPACK_PROBE_WINDOWS: usize = 512;
8
9#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, Default)]
11#[non_exhaustive]
12pub enum Format {
13 #[default]
15 Auto,
16 Mmcif,
18 Pdbml,
20 BinaryCif,
22 Mmtf,
24 Pdb,
26 Pqr,
28 Pdbqt,
30 Sdf,
32 Mol2,
34 SmallCif,
36}
37
38impl Format {
39 pub const NAMED: [Self; 10] = [
41 Self::Mmcif,
42 Self::Pdbml,
43 Self::BinaryCif,
44 Self::Mmtf,
45 Self::Pdbqt,
46 Self::Pqr,
47 Self::Pdb,
48 Self::Sdf,
49 Self::Mol2,
50 Self::SmallCif,
51 ];
52
53 #[must_use]
55 pub const fn name(self) -> &'static str {
56 match self {
57 Self::Auto => "auto",
58 Self::Mmcif => "mmcif",
59 Self::Pdbml => "pdbml",
60 Self::BinaryCif => "bcif",
61 Self::Mmtf => "mmtf",
62 Self::Pdb => "pdb",
63 Self::Pqr => "pqr",
64 Self::Pdbqt => "pdbqt",
65 Self::Sdf => "sdf",
66 Self::Mol2 => "mol2",
67 Self::SmallCif => "smallcif",
68 }
69 }
70
71 #[must_use]
73 pub const fn extensions(self) -> &'static [&'static str] {
74 match self {
75 Self::Auto | Self::SmallCif => &[],
76 Self::Mmcif => &["cif", "mmcif"],
77 Self::Pdbml => &["xml", "pdbml"],
78 Self::BinaryCif => &["bcif"],
79 Self::Mmtf => &["mmtf"],
80 Self::Pdb => &["pdb", "ent"],
81 Self::Pqr => &["pqr"],
82 Self::Pdbqt => &["pdbqt"],
83 Self::Sdf => &["sdf", "mol"],
84 Self::Mol2 => &["mol2"],
85 }
86 }
87
88 #[must_use]
90 pub fn parse(name: &str) -> Option<Self> {
91 if name.eq_ignore_ascii_case("auto") {
92 Some(Self::Auto)
93 } else if name.eq_ignore_ascii_case("mmcif") || name.eq_ignore_ascii_case("cif") {
94 Some(Self::Mmcif)
95 } else if name.eq_ignore_ascii_case("pdbml") || name.eq_ignore_ascii_case("xml") {
96 Some(Self::Pdbml)
97 } else if name.eq_ignore_ascii_case("bcif") || name.eq_ignore_ascii_case("binarycif") {
98 Some(Self::BinaryCif)
99 } else if name.eq_ignore_ascii_case("mmtf") {
100 Some(Self::Mmtf)
101 } else if name.eq_ignore_ascii_case("pdb") || name.eq_ignore_ascii_case("ent") {
102 Some(Self::Pdb)
103 } else if name.eq_ignore_ascii_case("pqr") {
104 Some(Self::Pqr)
105 } else if name.eq_ignore_ascii_case("pdbqt") {
106 Some(Self::Pdbqt)
107 } else if name.eq_ignore_ascii_case("sdf") || name.eq_ignore_ascii_case("mol") {
108 Some(Self::Sdf)
109 } else if name.eq_ignore_ascii_case("mol2") {
110 Some(Self::Mol2)
111 } else if name.eq_ignore_ascii_case("smallcif") {
112 Some(Self::SmallCif)
113 } else {
114 None
115 }
116 }
117
118 #[must_use]
120 pub fn recognises(self, bytes: &[u8]) -> bool {
121 match self {
122 Self::Auto | Self::Pqr => false,
123 Self::Mmcif => first_lines(bytes, 8).any(|line| line.starts_with(b"data_")),
124 Self::Pdbml => contains_in_probe(bytes, b"PDBx:datablock", PDBML_PROBE_WINDOWS),
125 Self::BinaryCif => recognises_messagepack_map(bytes, b"dataBlocks"),
126 Self::Mmtf => recognises_messagepack_map(bytes, b"mmtfVersion"),
127 Self::Pdbqt => recognises_pdbqt(bytes),
128 Self::Pdb => recognises_pdb(bytes),
129 Self::Sdf => recognises_sdf(bytes),
130 Self::Mol2 => recognises_mol2(bytes),
131 Self::SmallCif => recognises_small_cif(bytes),
132 }
133 }
134
135 pub fn detect(
143 requested: Self,
144 input: &InputBuffer,
145 name: Option<&str>,
146 ) -> Result<Self, Diagnostic> {
147 if requested != Self::Auto {
148 return Ok(requested);
149 }
150
151 let bytes = input.as_bytes();
152 for format in [
153 Self::SmallCif,
154 Self::Mmcif,
155 Self::Pdbml,
156 Self::BinaryCif,
157 Self::Mmtf,
158 Self::Pdbqt,
159 Self::Sdf,
160 Self::Mol2,
161 ] {
162 if format.recognises(bytes) {
163 return Ok(format);
164 }
165 }
166
167 let named = name.and_then(Self::from_name);
168 if recognises_pdb(bytes) {
169 if let Some(format @ (Self::Pqr | Self::Pdbqt)) = named {
170 return Ok(format);
171 }
172 return Ok(Self::Pdb);
173 }
174 if let Some(format) = named {
175 return Ok(format);
176 }
177
178 let Some(name) = name else {
179 return Err(Diagnostic::new(Code::E1001)
180 .with_context("tried", "content, then file name")
181 .with_context("name", "(none)"));
182 };
183 Err(Diagnostic::new(Code::E1001)
184 .with_context("tried", "content, then file name")
185 .with_context("name", name))
186 }
187
188 #[must_use]
190 pub fn from_name(name: &str) -> Option<Self> {
191 let mut suffixes = name.rsplit('.');
192 let mut candidate = suffixes.next()?;
193 if is_compression_suffix(candidate) {
194 candidate = suffixes.next()?;
195 }
196 Self::from_extension(candidate)
197 }
198
199 #[must_use]
202 pub fn from_extension(extension: &str) -> Option<Self> {
203 Self::NAMED.into_iter().find(|format| {
204 format
205 .extensions()
206 .iter()
207 .any(|known| extension.eq_ignore_ascii_case(known))
208 })
209 }
210}
211
212fn is_compression_suffix(suffix: &str) -> bool {
213 suffix.eq_ignore_ascii_case("gz")
214 || suffix.eq_ignore_ascii_case("zst")
215 || suffix.eq_ignore_ascii_case("bz2")
216 || suffix.eq_ignore_ascii_case("xz")
217}
218
219fn is_messagepack_map_prefix(byte: u8) -> bool {
220 (0x80..=0x8f).contains(&byte) || matches!(byte, 0xde | 0xdf)
221}
222
223fn recognises_messagepack_map(bytes: &[u8], key: &[u8]) -> bool {
224 bytes
225 .first()
226 .copied()
227 .is_some_and(is_messagepack_map_prefix)
228 && contains_in_probe(bytes, key, MESSAGEPACK_PROBE_WINDOWS)
229}
230
231fn contains_in_probe(bytes: &[u8], needle: &[u8], windows: usize) -> bool {
232 if needle.is_empty() {
233 return true;
234 }
235 bytes
236 .windows(needle.len())
237 .take(windows)
238 .any(|window| window == needle)
239}
240
241fn recognises_pdb(bytes: &[u8]) -> bool {
242 first_lines(bytes, 64).any(|line| {
243 line.starts_with(b"ATOM ")
244 || line.starts_with(b"HETATM")
245 || line.starts_with(b"HEADER")
246 || line.starts_with(b"CRYST1")
247 || line.starts_with(b"MODEL ")
248 })
249}
250
251fn recognises_pdbqt(bytes: &[u8]) -> bool {
252 first_lines(bytes, 64).any(recognises_pdbqt_line)
253}
254
255fn recognises_pdbqt_line(line: &[u8]) -> bool {
256 if matches!(line, b"ROOT" | b"ENDROOT") || line.starts_with(b"TORSDOF") {
257 return true;
258 }
259 if !(line.starts_with(b"ATOM ") || line.starts_with(b"HETATM")) {
260 return false;
261 }
262 let charge = line
263 .get(70..76)
264 .and_then(|field| str::from_utf8(field).ok());
265 let atom_type = line.get(77..).and_then(|field| str::from_utf8(field).ok());
266 charge.is_some_and(|field| field.trim().parse::<f64>().is_ok())
267 && atom_type.is_some_and(|field| !field.trim().is_empty())
268}
269
270fn recognises_sdf(bytes: &[u8]) -> bool {
271 let mut lines = first_lines(bytes, 64);
272 if lines.any(|line| line == b"$$$$") {
273 return true;
274 }
275 first_lines(bytes, 4)
276 .nth(3)
277 .is_some_and(|line| line.ends_with(b"V2000") || line.ends_with(b"V3000"))
278}
279
280fn recognises_mol2(bytes: &[u8]) -> bool {
281 first_lines(bytes, 64).any(|line| line.starts_with(b"@<TRIPOS>MOLECULE"))
282}
283
284fn recognises_small_cif(bytes: &[u8]) -> bool {
286 let mut data_block = false;
287 let mut fractional = false;
288 for line in first_lines(bytes, 256) {
289 data_block |= line.starts_with(b"data_");
290 fractional |=
291 line.starts_with(b"_atom_site_fract_x") || line.starts_with(b"_atom_site.fract_x");
292 if line.starts_with(b"_atom_site.Cartn_x") {
293 return false;
294 }
295 }
296 data_block && fractional
297}
298
299fn first_lines(bytes: &[u8], count: usize) -> impl Iterator<Item = &[u8]> {
300 bytes
301 .split(|byte| *byte == b'\n')
302 .take(count)
303 .map(strip_carriage_return)
304}
305
306fn strip_carriage_return(line: &[u8]) -> &[u8] {
307 if line.last().copied() == Some(b'\r') {
308 &line[..line.len() - 1]
309 } else {
310 line
311 }
312}