metatensor/block/
mod.rs

1// TensorBlock can be manipulated in three forms: as standalone block owning
2// it's data, or as a reference (or mutable reference) to a block inside a
3// `TensorMap`. The three forms are represented by a pointer to `mts_block_t`
4// in the C API, but we wrap them into three different types for Rust.
5
6mod block_ref;
7use crate::c_api::mts_block_t;
8
9pub use self::block_ref::{TensorBlockRef, TensorBlockData, GradientsIter};
10
11mod block_mut;
12pub use self::block_mut::{TensorBlockRefMut, TensorBlockDataMut, GradientsMutIter};
13
14mod owned;
15pub use self::owned::TensorBlock;
16
17
18
19/// Lazily accessed metadata inside a `TensorBlock`
20///
21/// This struct provides immutable access to an object of type `T`, that is
22/// conceptually a field of a `TensorBlock`. The object is initialized on the
23/// first access, potentially saving some computation/allocations if the object
24/// is not needed.
25pub struct LazyMetadata<T> {
26    block: *const mts_block_t,
27    init: fn(*const mts_block_t) -> T,
28    metadata: once_cell::sync::OnceCell<T>,
29}
30
31impl<T: std::fmt::Debug> std::fmt::Debug for LazyMetadata<T> {
32    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33        f.debug_struct("Lazy")
34            .field("block", &self.block)
35            .field("metadata", &self.metadata)
36            .finish_non_exhaustive()
37    }
38}
39
40impl<T> std::ops::Deref for LazyMetadata<T> {
41    type Target = T;
42
43    fn deref(&self) -> &Self::Target {
44        self.metadata.get_or_init(|| (self.init)(self.block))
45    }
46}
47
48impl<T> LazyMetadata<T> {
49    /// Create a new `LazyMetadata` with the given initialization function
50    pub fn new(init: fn(*const mts_block_t) -> T, block: *const mts_block_t) -> LazyMetadata<T> {
51        LazyMetadata {
52            block: block,
53            init: init,
54            metadata: once_cell::sync::OnceCell::new(),
55        }
56    }
57}