Skip to main content

omgkit_core/
view.rs

1//! 单分子的零拷贝视图。
2//!
3//! [`MolBatch`] 内部一律用全局下标,而算法和用户想要的是分子内局部下标。
4//! [`MolView`] 只持有 `&MolBatch` 和分子号,把两者之间的换算收在一处 ——
5//! 不复制任何数据,构造代价是几个字段的赋值。
6
7use crate::batch::MolBatch;
8use crate::builder::{AtomData, BondData, MolBuilder};
9
10/// 批中单个分子的零拷贝视图。
11///
12/// 所有下标参数与返回值都是**分子内局部**下标(0 起)。
13#[derive(Debug, Clone, Copy)]
14pub struct MolView<'a> {
15    batch: &'a MolBatch,
16    idx: u32,
17    atom_base: u32,
18    bond_base: u32,
19    n_atoms: u32,
20    n_bonds: u32,
21}
22
23impl<'a> MolView<'a> {
24    pub(crate) fn new(batch: &'a MolBatch, idx: u32) -> Self {
25        let i = idx as usize;
26        let atom_base = batch.mol_atom_offset[i];
27        let bond_base = batch.mol_bond_offset[i];
28        Self {
29            batch,
30            idx,
31            atom_base,
32            bond_base,
33            n_atoms: batch.mol_atom_offset[i + 1] - atom_base,
34            n_bonds: batch.mol_bond_offset[i + 1] - bond_base,
35        }
36    }
37
38    /// 该分子在批中的下标
39    #[must_use]
40    pub fn index(&self) -> u32 {
41        self.idx
42    }
43
44    /// 所属的批
45    #[must_use]
46    pub fn batch(&self) -> &'a MolBatch {
47        self.batch
48    }
49
50    /// 原子数
51    #[must_use]
52    pub fn num_atoms(&self) -> usize {
53        self.n_atoms as usize
54    }
55
56    /// 键数
57    #[must_use]
58    pub fn num_bonds(&self) -> usize {
59        self.n_bonds as usize
60    }
61
62    /// 是否为空分子
63    #[must_use]
64    pub fn is_empty(&self) -> bool {
65        self.n_atoms == 0
66    }
67
68    /// 分子名
69    #[must_use]
70    pub fn name(&self) -> Option<&'a str> {
71        self.batch.names[self.idx as usize].as_deref()
72    }
73
74    /// 局部原子下标 → 全局下标。不做越界检查。
75    #[must_use]
76    pub fn global_atom(&self, local: u32) -> u32 {
77        self.atom_base + local
78    }
79
80    /// 全局原子下标 → 局部下标。不属于本分子时返回 `None`。
81    #[must_use]
82    pub fn local_atom(&self, global: u32) -> Option<u32> {
83        global
84            .checked_sub(self.atom_base)
85            .filter(|&l| l < self.n_atoms)
86    }
87
88    /// 取原子。越界返回 `None`。
89    #[must_use]
90    pub fn atom(&self, local: u32) -> Option<AtomData> {
91        if local >= self.n_atoms {
92            return None;
93        }
94        let g = (self.atom_base + local) as usize;
95        let b = self.batch;
96        Some(AtomData {
97            atomic_num: b.atomic_num[g],
98            formal_charge: b.formal_charge[g],
99            isotope: b.isotope[g],
100            num_explicit_hs: b.num_explicit_hs[g],
101            num_implicit_hs: b.num_implicit_hs[g],
102            num_radical_electrons: b.num_radical_electrons[g],
103            atom_map: b.atom_map[g],
104            chiral_tag: b.chiral_tag[g],
105            stereo_perm: b.stereo_perm[g],
106            hybridization: b.hybridization[g],
107            flags: b.atom_flags[g],
108        })
109    }
110
111    /// 取键,端点为**局部**下标。越界返回 `None`。
112    #[must_use]
113    pub fn bond(&self, local: u32) -> Option<BondData> {
114        if local >= self.n_bonds {
115            return None;
116        }
117        let g = (self.bond_base + local) as usize;
118        let b = self.batch;
119        Some(BondData {
120            begin: b.bond_begin[g] - self.atom_base,
121            end: b.bond_end[g] - self.atom_base,
122            order: b.bond_order[g],
123            direction: b.bond_direction[g],
124            stereo: b.bond_stereo[g],
125            stereo_atoms: b.bond_stereo_atoms[g].map(|a| {
126                if a == BondData::NO_STEREO_ATOM {
127                    a
128                } else {
129                    a - self.atom_base
130                }
131            }),
132            flags: b.bond_flags[g],
133        })
134    }
135
136    /// 原子的度(不含隐式氢)
137    #[must_use]
138    pub fn degree(&self, local: u32) -> usize {
139        if local >= self.n_atoms {
140            return 0;
141        }
142        let g = (self.atom_base + local) as usize;
143        (self.batch.nbr_offset[g + 1] - self.batch.nbr_offset[g]) as usize
144    }
145
146    /// 遍历某原子的邻居,产出 `(邻居局部下标, 键局部下标)`。
147    ///
148    /// 顺序即键的插入顺序 —— 手性语义依赖于此,详见
149    /// [`batch`](crate::batch) 模块中 CSR 构建的说明。
150    pub fn neighbors(&self, local: u32) -> impl Iterator<Item = (u32, u32)> + 'a {
151        let (start, end) = if local < self.n_atoms {
152            let g = (self.atom_base + local) as usize;
153            (self.batch.nbr_offset[g], self.batch.nbr_offset[g + 1])
154        } else {
155            (0, 0)
156        };
157        let b = self.batch;
158        let (abase, bbase) = (self.atom_base, self.bond_base);
159        (start..end).map(move |k| {
160            let k = k as usize;
161            (b.nbr_atom[k] - abase, b.nbr_bond[k] - bbase)
162        })
163    }
164
165    /// 遍历全部原子,产出 `(局部下标, 原子)`。
166    pub fn atoms(&self) -> impl Iterator<Item = (u32, AtomData)> + '_ {
167        (0..self.n_atoms).map(move |i| (i, self.atom(i).expect("下标由 n_atoms 生成")))
168    }
169
170    /// 遍历全部键,产出 `(局部下标, 键)`。
171    pub fn bonds(&self) -> impl Iterator<Item = (u32, BondData)> + '_ {
172        (0..self.n_bonds).map(move |i| (i, self.bond(i).expect("下标由 n_bonds 生成")))
173    }
174
175    /// 拷贝回可变的 [`MolBuilder`],用于编辑或反应产物构建。
176    ///
177    /// 这是 `MolBatchBuilder::push` 的逆操作;二者的往返恒等性由测试保证。
178    #[must_use]
179    pub fn to_builder(&self) -> MolBuilder {
180        let mut m = MolBuilder::with_capacity(self.num_atoms(), self.num_bonds());
181        for (_, a) in self.atoms() {
182            m.add_atom_data(a);
183        }
184        for (_, bd) in self.bonds() {
185            m.add_bond_data(bd).expect("视图中的键端点必然合法");
186        }
187        if let Some(n) = self.name() {
188            m.set_name(n);
189        }
190        m
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197    use crate::batch::MolBatchBuilder;
198    use crate::types::{
199        AtomFlags, BondDirection, BondFlags, BondOrder, BondStereo, ChiralTag, Hybridization,
200    };
201
202    fn ethanol() -> MolBuilder {
203        let mut m = MolBuilder::new();
204        let c0 = m.add_atom(6);
205        let c1 = m.add_atom(6);
206        let o = m.add_atom(8);
207        m.add_bond(c0, c1, BondOrder::Single).unwrap();
208        m.add_bond(c1, o, BondOrder::Single).unwrap();
209        m.set_name("ethanol");
210        m
211    }
212
213    fn acetic_acid() -> MolBuilder {
214        // CC(=O)O
215        let mut m = MolBuilder::new();
216        let c0 = m.add_atom(6);
217        let c1 = m.add_atom(6);
218        let o1 = m.add_atom(8);
219        let o2 = m.add_atom(8);
220        m.add_bond(c0, c1, BondOrder::Single).unwrap();
221        m.add_bond(c1, o1, BondOrder::Double).unwrap();
222        m.add_bond(c1, o2, BondOrder::Single).unwrap();
223        m
224    }
225
226    fn batch_of(mols: &[MolBuilder]) -> MolBatch {
227        let mut bb = MolBatchBuilder::new();
228        for m in mols {
229            bb.push(m).unwrap();
230        }
231        bb.finish()
232    }
233
234    #[test]
235    fn local_indices_are_zero_based_per_molecule() {
236        let b = batch_of(&[ethanol(), acetic_acid()]);
237        let m1 = b.mol(1).unwrap();
238
239        assert_eq!(m1.num_atoms(), 4);
240        assert_eq!(m1.num_bonds(), 3);
241        // 第二个分子的原子在全局是 3..7,局部必须是 0..4
242        assert_eq!(m1.global_atom(0), 3);
243        assert_eq!(m1.local_atom(3), Some(0));
244        assert_eq!(m1.local_atom(2), None, "全局原子 2 属于第一个分子");
245        assert_eq!(m1.local_atom(7), None, "全局原子 7 已越过本分子");
246
247        let bond = m1.bond(1).unwrap();
248        assert_eq!((bond.begin, bond.end), (1, 2), "键端点应为局部下标");
249        assert_eq!(bond.order, BondOrder::Double);
250    }
251
252    #[test]
253    fn neighbors_are_local() {
254        let b = batch_of(&[ethanol(), acetic_acid()]);
255        let m1 = b.mol(1).unwrap();
256
257        // 乙酸的 C1 连着 C0、O1、O2
258        let mut nbrs: Vec<u32> = m1.neighbors(1).map(|(a, _)| a).collect();
259        nbrs.sort_unstable();
260        assert_eq!(nbrs, vec![0, 2, 3]);
261        assert_eq!(m1.degree(1), 3);
262        assert_eq!(m1.degree(0), 1);
263    }
264
265    #[test]
266    fn neighbor_bond_indices_are_local() {
267        let b = batch_of(&[ethanol(), acetic_acid()]);
268        let m1 = b.mol(1).unwrap();
269        for (_, bond_local) in m1.neighbors(1) {
270            assert!(
271                (bond_local as usize) < m1.num_bonds(),
272                "键下标 {bond_local} 未换算为局部"
273            );
274        }
275    }
276
277    #[test]
278    fn out_of_range_access_is_none() {
279        let b = batch_of(&[ethanol()]);
280        let m = b.mol(0).unwrap();
281        assert!(m.atom(3).is_none());
282        assert!(m.bond(2).is_none());
283        assert_eq!(m.degree(99), 0);
284        assert_eq!(m.neighbors(99).count(), 0);
285    }
286
287    /// 造一个**每个字段都取非默认值**的分子。
288    ///
289    /// 这里刻意用**完整的结构体字面量**(不写 `..Default::default()`):
290    /// 以后给 `AtomData` 或 `BondData` 加字段时,这里会直接编译失败,逼着人把
291    /// 新字段也纳入往返测试。
292    ///
293    /// 之所以要这道保险:漏同步列式存储不会报错,只会让某个字段在
294    /// builder → batch → builder 的往返中悄悄变回默认值,而这类 bug 要到很久
295    /// 以后才以"化学算错了"的形式冒出来。
296    fn every_field_set() -> MolBuilder {
297        let mut m = MolBuilder::new();
298        m.add_atom_data(AtomData {
299            atomic_num: 7,
300            formal_charge: -1,
301            isotope: 15,
302            num_explicit_hs: 2,
303            num_implicit_hs: 3,
304            num_radical_electrons: 1,
305            atom_map: 7,
306            chiral_tag: ChiralTag::Cw,
307            stereo_perm: 0,
308            hybridization: Hybridization::Sp3d2,
309            flags: AtomFlags::AROMATIC | AtomFlags::NO_IMPLICIT | AtomFlags::IN_RING,
310        });
311        m.add_atom_data(AtomData {
312            atomic_num: 16,
313            formal_charge: 2,
314            isotope: 34,
315            num_explicit_hs: 1,
316            num_implicit_hs: 0,
317            num_radical_electrons: 2,
318            atom_map: 3,
319            chiral_tag: ChiralTag::Octahedral,
320            stereo_perm: 25,
321            hybridization: Hybridization::Sp2d,
322            flags: AtomFlags::CONJUGATED,
323        });
324        m.add_bond_data(BondData {
325            begin: 0,
326            end: 1,
327            order: BondOrder::Dative,
328            direction: BondDirection::DownRight,
329            stereo: BondStereo::Trans,
330            // 故意不是 [0, 1]:参照原子要跟着基址平移,顺序反过来才能看出
331            // 平移是否作用在了正确的位置上
332            stereo_atoms: [1, 0],
333            flags: BondFlags::AROMATIC | BondFlags::IN_RING | BondFlags::CONJUGATED,
334        })
335        .expect("端点合法");
336        m.set_name("每字段非默认");
337        m
338    }
339
340    /// builder → batch → view → builder 必须恒等。
341    /// 这是 L0 最重要的性质:列式存储不能悄悄丢字段。
342    #[test]
343    fn roundtrip_through_batch_is_identity() {
344        let m = every_field_set();
345
346        // 夹在两个分子中间,确保偏移换算被真正考验
347        let b = batch_of(&[ethanol(), m.clone(), ethanol()]);
348        let back = b.mol(1).unwrap().to_builder();
349
350        assert_eq!(back.atoms(), m.atoms(), "原子列往返不一致");
351        assert_eq!(back.bonds(), m.bonds(), "键列往返不一致");
352        assert_eq!(back.name(), m.name());
353    }
354
355    /// 确认上面那个分子**真的**每个字段都偏离了默认值 —— 否则往返测试会
356    /// 在一个全是默认值的分子上"通过",什么也没验证。
357    #[test]
358    fn every_field_is_actually_non_default() {
359        let m = every_field_set();
360        let default_atom = AtomData::default();
361        for (i, a) in m.atoms().iter().enumerate() {
362            assert_ne!(a.atomic_num, default_atom.atomic_num, "原子{i}.元素");
363            assert_ne!(a.formal_charge, default_atom.formal_charge, "原子{i}.电荷");
364            assert_ne!(a.isotope, default_atom.isotope, "原子{i}.同位素");
365            assert_ne!(a.atom_map, default_atom.atom_map, "原子{i}.映射号");
366            assert_ne!(a.chiral_tag, default_atom.chiral_tag, "原子{i}.手性");
367            assert_ne!(a.hybridization, default_atom.hybridization, "原子{i}.杂化");
368            assert_ne!(a.flags, AtomFlags::NONE, "原子{i}.标志");
369        }
370        // 两个原子合起来覆盖显式氢/隐式氢/自由基的非零取值
371        assert!(m.atoms().iter().any(|a| a.num_explicit_hs != 0));
372        assert!(m.atoms().iter().any(|a| a.num_implicit_hs != 0));
373        assert!(m.atoms().iter().any(|a| a.num_radical_electrons != 0));
374        // 排列序号对四面体恒为 0,所以只能要求"至少一个原子非零",并且
375        // 两类立体标记(四面体 / 配位几何)都要出现
376        assert!(m.atoms().iter().any(|a| a.stereo_perm != 0));
377        assert!(m.atoms().iter().any(|a| a.chiral_tag.is_tetrahedral()));
378        assert!(m.atoms().iter().any(|a| !a.chiral_tag.is_tetrahedral()));
379
380        for (i, b) in m.bonds().iter().enumerate() {
381            assert_ne!(b.order, BondOrder::Unspecified, "键{i}.键级");
382            assert_ne!(b.direction, BondDirection::None, "键{i}.方向");
383            assert_ne!(b.stereo, BondStereo::None, "键{i}.立体");
384            // 参照原子存的是全局下标,取回来必须换算成局部 —— 少了换算,
385            // 批里第二个分子往后的参照就全指向别人家的原子
386            assert_eq!(b.stereo_atoms, [1, 0], "键{i}.立体参照原子");
387            assert_ne!(b.flags, BondFlags::NONE, "键{i}.标志");
388        }
389    }
390
391    #[test]
392    fn iter_visits_every_molecule() {
393        let b = batch_of(&[ethanol(), acetic_acid(), ethanol()]);
394        let sizes: Vec<usize> = b.iter().map(|m| m.num_atoms()).collect();
395        assert_eq!(sizes, vec![3, 4, 3]);
396        assert_eq!(b.iter().count(), 3);
397    }
398
399    #[test]
400    fn try_mol_reports_out_of_range() {
401        let b = batch_of(&[ethanol()]);
402        assert!(b.try_mol(0).is_ok());
403        let e = b.try_mol(5).unwrap_err();
404        assert!(matches!(
405            e,
406            crate::error::Error::MolIndexOutOfRange {
407                index: 5,
408                num_mols: 1
409            }
410        ));
411    }
412}