cuest_cudarc/
query.rs

1//! Typed queries of cuEST object attributes (`cuestQuery`).
2//!
3//! cuEST attributes are queried through one type-erased C function; the
4//! attribute determines the value type. This module encodes that mapping in
5//! the type system: each attribute is a marker type implementing
6//! [`CuestAttribute`] (which carries the value type and the C id), plus an
7//! object-scoped trait ([`AoshellAttribute`], [`AobasisAttribute`], ...) so
8//! that only valid attribute/object combinations compile.
9//!
10//! The v1 table covers the AO shell and AO basis attributes; more markers are
11//! pure transcription work. Unknown attributes can be reached through the raw
12//! bindings ([`cuest_ffi`]).
13
14use crate::error::{status_check, CuestError};
15use cuest_ffi::ffi;
16use std::ffi::{c_int, c_void};
17
18mod sealed {
19    pub trait Sealed {}
20    impl Sealed for u64 {}
21    impl Sealed for i32 {}
22    impl Sealed for f64 {}
23}
24
25/// Value types an attribute can be read into.
26///
27/// Sealed: the set of types cuEST returns for scalars is closed (`uint64_t`,
28/// `int32_t`, `double`).
29pub trait AttributeValue: Copy + sealed::Sealed {
30    /// Size in bytes cuEST expects for this type.
31    const SIZE: usize;
32    /// Interpret the value written by cuEST into `bytes`.
33    fn from_bytes(bytes: &[u8]) -> Self;
34}
35
36fn read_ne<const N: usize>(bytes: &[u8]) -> [u8; N] {
37    let mut buf = [0u8; N];
38    buf.copy_from_slice(&bytes[..N]);
39    buf
40}
41
42impl AttributeValue for u64 {
43    const SIZE: usize = size_of::<u64>();
44    fn from_bytes(bytes: &[u8]) -> Self {
45        u64::from_ne_bytes(read_ne(bytes))
46    }
47}
48
49impl AttributeValue for i32 {
50    const SIZE: usize = size_of::<i32>();
51    fn from_bytes(bytes: &[u8]) -> Self {
52        i32::from_ne_bytes(read_ne(bytes))
53    }
54}
55
56impl AttributeValue for f64 {
57    const SIZE: usize = size_of::<f64>();
58    fn from_bytes(bytes: &[u8]) -> Self {
59        f64::from_ne_bytes(read_ne(bytes))
60    }
61}
62
63/// A queryable attribute: its value type and its C attribute id.
64pub trait CuestAttribute {
65    /// The type cuEST returns for this attribute.
66    type Value: AttributeValue;
67    /// The `int attribute` id passed to `cuestQuery`.
68    const ID: c_int;
69}
70
71/// Attribute of a `cuestAOShell_t` object.
72pub trait AoshellAttribute: CuestAttribute {}
73
74/// Attribute of a `cuestAOBasis_t` object.
75pub trait AobasisAttribute: CuestAttribute {}
76
77macro_rules! attribute {
78    ($(#[$doc:meta])* $name:ident: $trait:ident, $value:ty, $id:expr) => {
79        $(#[$doc])*
80        pub struct $name;
81        impl CuestAttribute for $name {
82            type Value = $value;
83            const ID: c_int = $id;
84        }
85        impl $trait for $name {}
86    };
87}
88
89attribute!(
90    /// Whether the shell uses pure (spherical) functions (nonzero if true).
91    AoshellIsPure: AoshellAttribute, i32, ffi::CUEST_AOSHELL_IS_PURE.0 as c_int
92);
93attribute!(
94    /// Angular momentum quantum number (L).
95    AoshellL: AoshellAttribute, u64, ffi::CUEST_AOSHELL_L.0 as c_int
96);
97attribute!(
98    /// Number of primitive Gaussians.
99    AoshellNumPrimitive: AoshellAttribute, u64, ffi::CUEST_AOSHELL_NUM_PRIMITIVE.0 as c_int
100);
101attribute!(
102    /// Number of basis functions (AOs, either pure or Cartesian as set).
103    AoshellNumAo: AoshellAttribute, u64, ffi::CUEST_AOSHELL_NUM_AO.0 as c_int
104);
105attribute!(
106    /// Number of pure (spherical) AOs.
107    AoshellNumPure: AoshellAttribute, u64, ffi::CUEST_AOSHELL_NUM_PURE.0 as c_int
108);
109attribute!(
110    /// Number of Cartesian AOs.
111    AoshellNumCart: AoshellAttribute, u64, ffi::CUEST_AOSHELL_NUM_CART.0 as c_int
112);
113
114attribute!(
115    /// Nonzero if all shells are pure (spherical), zero otherwise.
116    AobasisIsPure: AobasisAttribute, i32, ffi::CUEST_AOBASIS_IS_PURE.0 as c_int
117);
118attribute!(
119    /// Nonzero if all shells are Cartesian, zero otherwise.
120    AobasisIsCart: AobasisAttribute, i32, ffi::CUEST_AOBASIS_IS_CART.0 as c_int
121);
122attribute!(
123    /// Nonzero if shells are mixed spherical/Cartesian, zero otherwise.
124    AobasisIsMixed: AobasisAttribute, i32, ffi::CUEST_AOBASIS_IS_MIXED.0 as c_int
125);
126attribute!(
127    /// Number of basis function centers.
128    AobasisNumAtom: AobasisAttribute, u64, ffi::CUEST_AOBASIS_NUM_ATOM.0 as c_int
129);
130attribute!(
131    /// Total number of AO shells.
132    AobasisNumShell: AobasisAttribute, u64, ffi::CUEST_AOBASIS_NUM_SHELL.0 as c_int
133);
134attribute!(
135    /// Total number of basis functions (AOs).
136    AobasisNumAo: AobasisAttribute, u64, ffi::CUEST_AOBASIS_NUM_AO.0 as c_int
137);
138attribute!(
139    /// Total number of Cartesian basis functions.
140    AobasisNumCart: AobasisAttribute, u64, ffi::CUEST_AOBASIS_NUM_CART.0 as c_int
141);
142attribute!(
143    /// Total number of Gaussian primitives.
144    AobasisNumPrimitive: AobasisAttribute, u64, ffi::CUEST_AOBASIS_NUM_PRIMITIVE.0 as c_int
145);
146attribute!(
147    /// Maximum angular momentum (L) among shells.
148    AobasisMaxL: AobasisAttribute, u64, ffi::CUEST_AOBASIS_MAX_L.0 as c_int
149);
150
151/// Run a `cuestQuery` for a scalar attribute of an object.
152pub(crate) fn query_attribute<V: AttributeValue>(
153    handle: ffi::cuestHandle_t,
154    object_type: ffi::cuestType_t,
155    object: *const c_void,
156    attribute: c_int,
157) -> Result<V, CuestError> {
158    let mut buf = [0u8; size_of::<u64>()];
159    status_check(
160        unsafe {
161            ffi::cuestQuery(
162                handle,
163                object_type,
164                object,
165                attribute,
166                buf.as_mut_ptr() as *mut c_void,
167                V::SIZE,
168            )
169        },
170        "cuestQuery",
171    )?;
172    Ok(V::from_bytes(&buf))
173}
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178
179    #[test]
180    fn cpu_attribute_ids_match_c() {
181        assert_eq!(<AoshellIsPure as CuestAttribute>::ID, 0);
182        assert_eq!(<AoshellNumCart as CuestAttribute>::ID, 5);
183        assert_eq!(<AobasisIsPure as CuestAttribute>::ID, 0);
184        assert_eq!(<AobasisMaxL as CuestAttribute>::ID, 8);
185    }
186
187    #[test]
188    fn cpu_value_sizes() {
189        assert_eq!(<u64 as AttributeValue>::SIZE, 8);
190        assert_eq!(<i32 as AttributeValue>::SIZE, 4);
191        assert_eq!(<f64 as AttributeValue>::SIZE, 8);
192    }
193
194    #[test]
195    fn cpu_from_bytes() {
196        assert_eq!(u64::from_bytes(&7u64.to_ne_bytes()), 7);
197        assert_eq!(i32::from_bytes(&(-3i32).to_ne_bytes()), -3);
198        assert_eq!(f64::from_bytes(&1.5f64.to_ne_bytes()), 1.5);
199        // Zero-padded tail (the query buffer is wider than the value).
200        assert_eq!(i32::from_bytes(&[2, 0, 0, 0, 9, 9, 9, 9]), 2);
201    }
202}
cuest_cudarc/query.rs

cuest_cudarc/
query.rs
