Skip to main content

tenferro_tensor/storage/
group.rs

1use std::fmt;
2use std::marker::PhantomData;
3use std::mem::{align_of, size_of};
4use std::ptr::NonNull;
5
6use crate::types::{tensor_from_group, tensor_view_from_group};
7use crate::{DType, DynRank, Placement, TensorLayout, TensorRank, TensorRead, TensorScalar};
8use smallvec::SmallVec;
9
10use super::prepared::{
11    prepare_read, prepare_write, validate_descriptor, AccessError, AccessTarget, CheckedDescriptor,
12    CheckedRead, CheckedWrite, PreparedRead, PreparedWrite, ProviderReadMapping,
13    ProviderWriteMapping, WriteInjectivityProof,
14};
15use super::root::{BackendAllocation, OwnedStorage, ProviderKind};
16use super::span::{ByteRange, RootBoundSpan};
17
18/// A group-local append-only allocation entry.
19#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
20pub(crate) struct AllocationSlot(u32);
21
22impl AllocationSlot {
23    pub(crate) const fn index(self) -> usize {
24        self.0 as usize
25    }
26
27    #[cfg(test)]
28    pub(crate) const fn test_raw(raw: u32) -> Self {
29        Self(raw)
30    }
31}
32
33/// A group-local descriptor lookup key. It carries no ownership authority.
34///
35/// # Examples
36///
37/// ```
38/// use tenferro_tensor::DescriptorSlot;
39///
40/// let slot = DescriptorSlot::from_index(3).unwrap();
41/// assert_eq!(slot.index(), 3);
42/// ```
43#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
44pub struct DescriptorSlot(u32);
45
46impl DescriptorSlot {
47    /// Return the zero-based descriptor index.
48    ///
49    /// # Examples
50    ///
51    /// ```
52    /// use tenferro_tensor::DescriptorSlot;
53    ///
54    /// let slot = DescriptorSlot::from_index(2).unwrap();
55    /// assert_eq!(slot.index(), 2);
56    /// ```
57    pub const fn index(self) -> usize {
58        self.0 as usize
59    }
60
61    /// Convert a host index into a descriptor slot.
62    ///
63    /// # Examples
64    ///
65    /// ```
66    /// use tenferro_tensor::DescriptorSlot;
67    ///
68    /// assert_eq!(DescriptorSlot::from_index(1).unwrap().index(), 1);
69    /// assert!(DescriptorSlot::from_index(usize::MAX).is_none());
70    /// ```
71    pub fn from_index(index: usize) -> Option<Self> {
72        match u32::try_from(index) {
73            Ok(index) => Some(Self(index)),
74            Err(_) => None,
75        }
76    }
77
78    #[cfg(test)]
79    pub(crate) const fn test_raw(raw: u32) -> Self {
80        Self(raw)
81    }
82}
83
84/// Construction input for one rank-specific descriptor.
85#[derive(Clone, Debug, PartialEq, Eq)]
86pub(crate) struct DescriptorInput<R: TensorRank> {
87    relative: ByteRange,
88    shape: R::Shape,
89    strides: R::Strides,
90    offset: isize,
91    require_injective: bool,
92}
93
94impl<R: TensorRank> DescriptorInput<R> {
95    pub(crate) fn new(
96        relative: ByteRange,
97        shape: R::Shape,
98        strides: R::Strides,
99        offset: isize,
100        require_injective: bool,
101    ) -> Self {
102        Self {
103            relative,
104            shape,
105            strides,
106            offset,
107            require_injective,
108        }
109    }
110}
111
112/// One validated, non-owning logical descriptor.
113#[derive(Clone, Debug, PartialEq, Eq)]
114pub(crate) struct DescriptorRecord {
115    allocation: AllocationSlot,
116    root: super::identity::RootResourceIdentity,
117    span: RootBoundSpan,
118    layout: TensorLayout<DynRank>,
119    dtype: DType,
120    element_size: usize,
121    element_count: usize,
122    provider: ProviderKind,
123    placement: Placement,
124    envelope: Option<ByteRange>,
125    write_injective: bool,
126    checked: CheckedDescriptor<DynRank>,
127}
128
129impl DescriptorRecord {
130    pub(crate) const fn allocation(&self) -> AllocationSlot {
131        self.allocation
132    }
133
134    pub(crate) const fn span(&self) -> RootBoundSpan {
135        self.span
136    }
137
138    pub(crate) const fn dtype(&self) -> DType {
139        self.dtype
140    }
141
142    pub(crate) const fn element_size(&self) -> usize {
143        self.element_size
144    }
145
146    pub(crate) const fn element_count(&self) -> usize {
147        self.element_count
148    }
149
150    pub(crate) fn layout(&self) -> &TensorLayout<DynRank> {
151        &self.layout
152    }
153
154    pub(crate) const fn provider(&self) -> ProviderKind {
155        self.provider
156    }
157
158    pub(crate) fn placement(&self) -> &Placement {
159        &self.placement
160    }
161
162    pub(crate) const fn envelope(&self) -> Option<ByteRange> {
163        self.envelope
164    }
165
166    pub(crate) const fn write_injective(&self) -> bool {
167        self.write_injective
168    }
169}
170
171/// Group construction and slot errors.
172///
173/// # Examples
174///
175/// ```
176/// use tenferro_tensor::GroupError;
177///
178/// let error = GroupError::IndexOverflow;
179/// assert!(error.to_string().contains("overflows"));
180/// ```
181#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
182pub enum GroupError {
183    #[error("group index overflows u32")]
184    IndexOverflow,
185    #[error("descriptor slot {slot} is outside the group")]
186    DescriptorSlotOutOfBounds { slot: usize },
187    #[error("descriptor slot {slot} is vacant")]
188    DescriptorSlotVacant { slot: usize },
189    #[error("allocation slot {slot} is outside the group")]
190    AllocationSlotOutOfBounds { slot: usize },
191    #[error("allocation slot {slot} is vacant")]
192    AllocationSlotVacant { slot: usize },
193    #[error("descriptor validation failed: {message}")]
194    InvalidDescriptor { message: String },
195    #[error("descriptor dtype mismatch: expected {expected:?}, actual {actual:?}")]
196    DTypeMismatch { expected: DType, actual: DType },
197    #[error("descriptor rank mismatch: expected {expected}, actual {actual}")]
198    RankMismatch { expected: usize, actual: usize },
199    #[error("allocation slot {allocation} has more than one live descriptor")]
200    AliasedAllocation { allocation: usize },
201}
202
203/// N-way mutable split errors. Every error leaves the group unchanged.
204#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
205pub(crate) enum DisjointViewError {
206    #[error(transparent)]
207    Group(#[from] GroupError),
208    #[error("descriptor slot {slot} appears more than once")]
209    DuplicateSlot { slot: usize },
210    #[error("descriptor slot {slot} has a non-injective mutable layout")]
211    NonInjective { slot: usize },
212    #[error("requested mutable descriptor envelopes overlap")]
213    PairwiseOverlap,
214    #[error("requested mutable descriptors are not provably disjoint")]
215    NotProvablyDisjoint,
216}
217
218/// Structural extraction errors.
219#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
220pub(crate) enum ExtractError {
221    #[error(transparent)]
222    Group(#[from] GroupError),
223    #[error("allocation slot {allocation} still has another descriptor")]
224    AliasedAllocation { allocation: usize },
225}
226
227/// One group of move-only owners and append-only logical descriptors.
228///
229/// # Examples
230///
231/// ```
232/// use tenferro_tensor::AllocationGroup;
233///
234/// let (group, bindings) = AllocationGroup::from_tensors(Vec::new())?;
235/// assert!(bindings.is_empty());
236/// assert!(format!("{group:?}").contains("AllocationGroup"));
237/// # Ok::<(), tenferro_tensor::GroupError>(())
238/// ```
239#[derive(Default)]
240pub struct AllocationGroup {
241    // Most public tensors contain one root and one descriptor. Keep that
242    // common case inline so the ownership boundary does not add a per-result
243    // metadata allocation to CPU hot paths; the vectors still grow for
244    // explicit multi-descriptor groups.
245    allocations: SmallVec<[Option<OwnedStorage>; 1]>,
246    descriptors: SmallVec<[Option<DescriptorRecord>; 1]>,
247}
248
249/// A shared descriptor child bounded by the group's shared borrow.
250pub(crate) struct GroupReadView<'a, T, R: TensorRank> {
251    owner: NonNull<OwnedStorage>,
252    descriptor: DescriptorRecord,
253    _borrow: PhantomData<(&'a OwnedStorage, T, R)>,
254}
255
256// SAFETY: the view is a shared capability over an immutable root borrow;
257// `OwnedStorage` is `Sync`, and the descriptor contains only checked metadata.
258unsafe impl<'a, T: Send, R: TensorRank> Send for GroupReadView<'a, T, R> {}
259unsafe impl<'a, T: Sync, R: TensorRank> Sync for GroupReadView<'a, T, R> {}
260
261impl<'a, T, R: TensorRank> Clone for GroupReadView<'a, T, R> {
262    fn clone(&self) -> Self {
263        Self {
264            owner: self.owner,
265            descriptor: self.descriptor.clone(),
266            _borrow: PhantomData,
267        }
268    }
269}
270
271impl<'a, T, R: TensorRank> GroupReadView<'a, T, R> {
272    pub(crate) fn clone_dyn(&self) -> GroupReadView<'a, T, crate::DynRank> {
273        GroupReadView {
274            owner: self.owner,
275            descriptor: self.descriptor.clone(),
276            _borrow: PhantomData,
277        }
278    }
279}
280
281impl<T, R: TensorRank> std::fmt::Debug for GroupReadView<'_, T, R> {
282    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
283        formatter
284            .debug_struct("GroupReadView")
285            .field("descriptor", &self.descriptor)
286            .finish_non_exhaustive()
287    }
288}
289
290impl<'a, T: TensorScalar, R: TensorRank> GroupReadView<'a, T, R> {
291    pub(crate) fn descriptor(&self) -> &DescriptorRecord {
292        &self.descriptor
293    }
294
295    pub(crate) fn provider_kind(&self) -> ProviderKind {
296        self.descriptor.provider
297    }
298
299    pub(crate) fn backend_identity(
300        &self,
301    ) -> Option<(crate::AllocationDomainId, crate::AllocationId)> {
302        if self.descriptor.provider == ProviderKind::Cpu {
303            return None;
304        }
305        let key = unsafe { self.owner.as_ref().root_identity().extent().key() };
306        Some((key.domain(), key.local()))
307    }
308
309    pub(crate) fn storage_buffer(&self) -> Option<&'a crate::StorageBuffer<T>> {
310        // SAFETY: the owner pointer is bounded by the group borrow carried by
311        // this view, and root storage never changes its concrete buffer.
312        let buffer = unsafe {
313            self.owner
314                .as_ref()
315                .host_buffer::<T>()
316                .or_else(|| self.owner.as_ref().backend_buffer::<T>())
317        }?;
318        Some(unsafe {
319            std::mem::transmute::<&crate::StorageBuffer<T>, &'a crate::StorageBuffer<T>>(buffer)
320        })
321    }
322
323    pub(crate) fn map_read(&self) -> Result<ProviderReadMapping<'_>, AccessError> {
324        // SAFETY: `owner` points into the group borrowed for `'a`; this method
325        // creates only a shared owner capability and maps the retained span.
326        unsafe {
327            self.owner
328                .as_ref()
329                .as_ref()
330                .map_read(self.descriptor.span, self.descriptor.dtype)
331        }
332    }
333
334    pub(crate) fn backend_allocation(&self) -> Option<&'a dyn BackendAllocation> {
335        let allocation = unsafe { self.owner.as_ref().backend_allocation() }?;
336        Some(unsafe {
337            std::mem::transmute::<&dyn BackendAllocation, &'a dyn BackendAllocation>(allocation)
338        })
339    }
340
341    pub(crate) fn prepare_device_read_for_layout(
342        &self,
343        layout: &TensorLayout<R>,
344    ) -> Result<Box<dyn crate::PreparedDeviceAccess + 'a>, AccessError> {
345        let owner: crate::storage::root::StorageRef<'a> = unsafe { self.owner.as_ref().as_ref() };
346        let checked: CheckedRead<'a, R> = CheckedRead::new::<T>(
347            // SAFETY: `owner` is bounded by the group's shared borrow.
348            owner,
349            self.descriptor.span,
350            R::shape_from_vec(layout.shape().iter().copied().collect()).map_err(|error| {
351                AccessError::InvalidLayout {
352                    message: error.to_string(),
353                }
354            })?,
355            R::strides_from_vec(layout.strides().iter().copied().collect()).map_err(|error| {
356                AccessError::InvalidLayout {
357                    message: error.to_string(),
358                }
359            })?,
360            layout.offset(),
361        )?;
362        prepare_read::<T, R>(checked, AccessTarget::Device)
363            .map_err(|failure| failure.1)?
364            .into_device_state()
365    }
366
367    pub(crate) fn prepare_host_read(&self) -> Result<PreparedRead<'_, T, DynRank>, AccessError> {
368        prepare_read(
369            CheckedRead::from_validated(
370                // SAFETY: the pointer is bounded by the group's shared borrow.
371                unsafe { self.owner.as_ref().as_ref() },
372                self.descriptor.checked.clone(),
373            ),
374            AccessTarget::Host,
375        )
376        .map_err(|failure| failure.1)
377    }
378
379    pub(crate) fn host_slice(&self) -> Result<&'a [T], AccessError> {
380        // SAFETY: `owner` is bounded by the group borrow carried by `'a`.
381        unsafe {
382            self.owner
383                .as_ref()
384                .as_ref()
385                .host_slice(self.descriptor.span, self.descriptor.dtype)
386        }
387    }
388}
389
390impl<'a, T: 'static, R: TensorRank> GroupReadView<'a, T, R> {
391    pub(crate) fn backend_buffer(&self) -> Option<&'a crate::StorageBuffer<T>> {
392        // SAFETY: the owner pointer is bounded by the group borrow carried by
393        // this view, and the root buffer cannot be resized after import.
394        let buffer = unsafe { self.owner.as_ref().backend_buffer::<T>() }?;
395        Some(unsafe {
396            std::mem::transmute::<&crate::StorageBuffer<T>, &'a crate::StorageBuffer<T>>(buffer)
397        })
398    }
399}
400
401/// A non-cloneable mutable descriptor child bounded by the group's exclusive
402/// borrow. The raw owner pointer is never exposed and is dereferenced only for
403/// a provider mapping whose retained byte envelope was proven by the group.
404pub(crate) struct GroupWriteView<'a, T, R: TensorRank> {
405    owner: NonNull<OwnedStorage>,
406    descriptor: DescriptorRecord,
407    _borrow: PhantomData<(&'a mut [u8], T, R)>,
408}
409
410// SAFETY: this is the exclusive capability for one owner borrow; moving it
411// transfers that exclusive borrow and cannot create a second access path.
412unsafe impl<'a, T: Send, R: TensorRank> Send for GroupWriteView<'a, T, R> {}
413
414impl<T, R: TensorRank> std::fmt::Debug for GroupWriteView<'_, T, R> {
415    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
416        formatter
417            .debug_struct("GroupWriteView")
418            .field("descriptor", &self.descriptor)
419            .finish_non_exhaustive()
420    }
421}
422
423impl<'a, T: TensorScalar, R: TensorRank> GroupWriteView<'a, T, R> {
424    pub(crate) fn descriptor(&self) -> &DescriptorRecord {
425        &self.descriptor
426    }
427
428    pub(crate) fn map_write(&mut self) -> Result<ProviderWriteMapping<'_>, AccessError> {
429        // SAFETY: the group proof established that this child owns a distinct
430        // reachable byte envelope; the temporary exclusive capability does not
431        // escape this mapping call except through that provider span guard.
432        unsafe {
433            self.owner
434                .as_mut()
435                .as_mut()
436                .map_write(self.descriptor.span, self.descriptor.dtype)
437        }
438    }
439
440    pub(crate) fn backend_buffer_mut(&mut self) -> Option<&'a mut crate::StorageBuffer<T>> {
441        // SAFETY: this child carries the group's exclusive borrow, so the
442        // mutable root projection cannot alias another owner capability.
443        let owner = unsafe { &mut *self.owner.as_ptr() };
444        let buffer = owner.backend_buffer_mut::<T>()?;
445        Some(unsafe {
446            std::mem::transmute::<&mut crate::StorageBuffer<T>, &'a mut crate::StorageBuffer<T>>(
447                buffer,
448            )
449        })
450    }
451
452    pub(crate) fn prepare_device_write_for_layout(
453        &mut self,
454        layout: &TensorLayout<R>,
455    ) -> Result<Box<dyn crate::PreparedDeviceAccess + 'a>, AccessError> {
456        let owner: crate::storage::root::StorageMut<'a> =
457            unsafe { (&mut *self.owner.as_ptr()).as_mut() };
458        let checked: CheckedWrite<'a, R> = CheckedWrite::new::<T>(
459            // SAFETY: this child carries the group's exclusive borrow.
460            owner,
461            self.descriptor.span,
462            R::shape_from_vec(layout.shape().iter().copied().collect()).map_err(|error| {
463                AccessError::InvalidLayout {
464                    message: error.to_string(),
465                }
466            })?,
467            R::strides_from_vec(layout.strides().iter().copied().collect()).map_err(|error| {
468                AccessError::InvalidLayout {
469                    message: error.to_string(),
470                }
471            })?,
472            layout.offset(),
473        )?;
474        prepare_write::<T, R>(checked, AccessTarget::Device)
475            .map_err(|failure| failure.1)?
476            .into_device_state()
477    }
478
479    pub(crate) fn prepare_host_write(
480        &mut self,
481    ) -> Result<PreparedWrite<'_, T, DynRank>, AccessError> {
482        let checked = CheckedWrite::from_validated(
483            // SAFETY: the pointer is bounded by this child group's exclusive borrow.
484            unsafe { self.owner.as_mut().as_mut() },
485            self.descriptor.checked.clone(),
486            WriteInjectivityProof,
487        );
488        prepare_write(checked, AccessTarget::Host).map_err(|failure| failure.1)
489    }
490
491    pub(crate) fn host_slice_mut(&mut self) -> Result<&'a mut [T], AccessError> {
492        // SAFETY: the group borrow is exclusive for `'a` and the descriptor
493        // retains the checked span used by the private host root.
494        unsafe {
495            self.owner
496                .as_mut()
497                .as_mut()
498                .host_slice_mut(self.descriptor.span, self.descriptor.dtype)
499        }
500    }
501}
502
503impl<'a, T: 'static, R: TensorRank> GroupWriteView<'a, T, R> {
504    pub(crate) fn backend_buffer(&self) -> Option<&crate::StorageBuffer<T>> {
505        // SAFETY: the owner pointer is bounded by the group's exclusive borrow;
506        // this shared inspection does not expose a mutable projection.
507        unsafe { self.owner.as_ref().backend_buffer::<T>() }
508    }
509}
510
511impl fmt::Debug for AllocationGroup {
512    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
513        formatter
514            .debug_struct("AllocationGroup")
515            .field("allocation_count", &self.allocations.len())
516            .field("descriptor_count", &self.descriptors.len())
517            .finish()
518    }
519}
520
521impl AllocationGroup {
522    pub(crate) fn new() -> Self {
523        Self::default()
524    }
525
526    /// Build one move-only group for detached runtime input ownership.
527    ///
528    /// # Examples
529    ///
530    /// ```
531    /// use tenferro_tensor::AllocationGroup;
532    ///
533    /// let (group, bindings) = AllocationGroup::from_tensors(Vec::new())?;
534    /// assert!(bindings.is_empty());
535    /// assert!(format!("{group:?}").contains("AllocationGroup"));
536    /// # Ok::<(), tenferro_tensor::GroupError>(())
537    /// ```
538    ///
539    /// # Errors
540    ///
541    /// Returns [`GroupError::IndexOverflow`] when the input count cannot be
542    /// represented by a descriptor slot.
543    pub fn from_tensors(
544        tensors: Vec<crate::Tensor>,
545    ) -> Result<(Self, Box<[DescriptorSlot]>), GroupError> {
546        let mut group = Self::new();
547        let mut bindings = Vec::with_capacity(tensors.len());
548        for tensor in tensors {
549            let (source, source_slot) = tensor.into_group_parts();
550            bindings.push(group.append_group(source, source_slot)?);
551        }
552        Ok((group, bindings.into_boxed_slice()))
553    }
554
555    /// Borrow dtype-erased read views for descriptor bindings without
556    /// materializing or cloning any owner.
557    ///
558    /// # Errors
559    ///
560    /// Returns [`GroupError::DescriptorSlotOutOfBounds`],
561    /// [`GroupError::DescriptorSlotVacant`], or [`GroupError::InvalidDescriptor`]
562    /// when a binding does not identify a valid descriptor.
563    pub fn read_views<'a>(
564        &'a self,
565        bindings: &[DescriptorSlot],
566    ) -> Result<Vec<TensorRead<'a>>, GroupError> {
567        bindings
568            .iter()
569            .map(|&slot| self.tensor_read(slot))
570            .collect()
571    }
572
573    /// Borrow one dtype-erased read view for a descriptor without materializing
574    /// or cloning its physical owner.
575    ///
576    /// # Errors
577    ///
578    /// Returns [`GroupError::DescriptorSlotOutOfBounds`],
579    /// [`GroupError::DescriptorSlotVacant`], or [`GroupError::InvalidDescriptor`]
580    /// when `slot` is not a valid descriptor.
581    pub fn read_view<'a>(&'a self, slot: DescriptorSlot) -> Result<TensorRead<'a>, GroupError> {
582        self.tensor_read(slot)
583    }
584
585    fn tensor_read<'a>(&'a self, slot: DescriptorSlot) -> Result<TensorRead<'a>, GroupError> {
586        let dtype = self.resolve_descriptor(slot)?.1.dtype();
587        let view = match dtype {
588            DType::F32 => tensor_view_from_group(self.view::<f32, DynRank>(slot)?),
589            DType::F64 => tensor_view_from_group(self.view::<f64, DynRank>(slot)?),
590            DType::I32 => tensor_view_from_group(self.view::<i32, DynRank>(slot)?),
591            DType::I64 => tensor_view_from_group(self.view::<i64, DynRank>(slot)?),
592            DType::Bool => tensor_view_from_group(self.view::<bool, DynRank>(slot)?),
593            DType::C32 => {
594                tensor_view_from_group(self.view::<num_complex::Complex32, DynRank>(slot)?)
595            }
596            DType::C64 => {
597                tensor_view_from_group(self.view::<num_complex::Complex64, DynRank>(slot)?)
598            }
599        }
600        .map_err(|error| GroupError::InvalidDescriptor {
601            message: error.to_string(),
602        })?;
603        Ok(TensorRead::from_view(view))
604    }
605
606    /// Append a tensor owner and return its new descriptor slot without copying.
607    ///
608    /// # Errors
609    ///
610    /// Returns [`GroupError::IndexOverflow`] when allocation or descriptor
611    /// indices cannot be represented, or [`GroupError::InvalidDescriptor`]
612    /// when the consumed tensor descriptor is invalid.
613    pub fn append_tensor(&mut self, tensor: crate::Tensor) -> Result<DescriptorSlot, GroupError> {
614        let (source, source_slot) = tensor.into_group_parts();
615        self.append_group(source, source_slot)
616    }
617
618    /// Append one descriptor and all of its physical owners without copying.
619    ///
620    /// # Errors
621    ///
622    /// Returns [`GroupError::IndexOverflow`] when group indices overflow or
623    /// [`GroupError::InvalidDescriptor`] when `source_slot` is invalid.
624    pub fn append_group(
625        &mut self,
626        mut source: AllocationGroup,
627        source_slot: DescriptorSlot,
628    ) -> Result<DescriptorSlot, GroupError> {
629        let allocation_offset =
630            u32::try_from(self.allocations.len()).map_err(|_| GroupError::IndexOverflow)?;
631        let descriptor_offset =
632            u32::try_from(self.descriptors.len()).map_err(|_| GroupError::IndexOverflow)?;
633        source.resolve_descriptor(source_slot)?;
634
635        for owner in source.allocations.drain(..) {
636            self.allocations.push(owner);
637        }
638        for descriptor in source.descriptors.drain(..) {
639            let descriptor = match descriptor {
640                Some(mut descriptor) => {
641                    let allocation = descriptor
642                        .allocation
643                        .0
644                        .checked_add(allocation_offset)
645                        .ok_or(GroupError::IndexOverflow)?;
646                    descriptor.allocation = AllocationSlot(allocation);
647                    Some(descriptor)
648                }
649                None => None,
650            };
651            self.descriptors.push(descriptor);
652        }
653
654        let source_descriptor_index =
655            u32::try_from(source_slot.index()).map_err(|_| GroupError::IndexOverflow)?;
656        Ok(DescriptorSlot(
657            source_descriptor_index
658                .checked_add(descriptor_offset)
659                .ok_or(GroupError::IndexOverflow)?,
660        ))
661    }
662
663    pub(crate) fn set_descriptor_placement(
664        &mut self,
665        slot: DescriptorSlot,
666        placement: Placement,
667    ) -> Result<(), GroupError> {
668        let descriptor = self
669            .descriptors
670            .get_mut(slot.index())
671            .ok_or(GroupError::DescriptorSlotOutOfBounds { slot: slot.index() })?
672            .as_mut()
673            .ok_or(GroupError::DescriptorSlotVacant { slot: slot.index() })?;
674        descriptor.placement = placement;
675        Ok(())
676    }
677
678    pub(crate) fn from_host_vec<T: TensorScalar, R: TensorRank>(
679        shape: R::Shape,
680        data: Vec<T>,
681    ) -> Result<(Self, DescriptorSlot), GroupError> {
682        let owner =
683            super::root::import_host_vec(data).map_err(|error| GroupError::InvalidDescriptor {
684                message: error.to_string(),
685            })?;
686        let span = owner.root_span();
687        let mut group = Self::new();
688        let allocation = group.insert_owner(owner)?;
689        let layout =
690            TensorLayout::<R>::compact(shape).map_err(|error| GroupError::InvalidDescriptor {
691                message: error.to_string(),
692            })?;
693        let input = DescriptorInput::new(
694            ByteRange::new(0, span.byte_len()),
695            R::shape_from_vec(layout.shape().iter().copied().collect()).map_err(|error| {
696                GroupError::InvalidDescriptor {
697                    message: error.to_string(),
698                }
699            })?,
700            R::strides_from_vec(layout.strides().iter().copied().collect()).map_err(|error| {
701                GroupError::InvalidDescriptor {
702                    message: error.to_string(),
703                }
704            })?,
705            layout.offset(),
706            true,
707        );
708        let slot = group.insert_descriptor::<T, R>(allocation, input)?;
709        Ok((group, slot))
710    }
711
712    /// Build one compact descriptor by consuming a scalar-independent provider root.
713    #[doc(hidden)]
714    pub fn from_backend_allocation<T: TensorScalar, R: TensorRank>(
715        shape: R::Shape,
716        allocation: Box<dyn BackendAllocation>,
717    ) -> Result<(Self, DescriptorSlot), GroupError> {
718        let owner = super::root::import_unique_root(allocation).map_err(|error| {
719            GroupError::InvalidDescriptor {
720                message: error.to_string(),
721            }
722        })?;
723        let span = owner.root_span();
724        let mut group = Self::new();
725        let allocation = group.insert_owner(owner)?;
726        let layout =
727            TensorLayout::<R>::compact(shape).map_err(|error| GroupError::InvalidDescriptor {
728                message: error.to_string(),
729            })?;
730        let input = DescriptorInput::new(
731            ByteRange::new(0, span.byte_len()),
732            R::shape_from_vec(layout.shape().iter().copied().collect()).map_err(|error| {
733                GroupError::InvalidDescriptor {
734                    message: error.to_string(),
735                }
736            })?,
737            R::strides_from_vec(layout.strides().iter().copied().collect()).map_err(|error| {
738                GroupError::InvalidDescriptor {
739                    message: error.to_string(),
740                }
741            })?,
742            layout.offset(),
743            true,
744        );
745        let slot = group.insert_descriptor::<T, R>(allocation, input)?;
746        Ok((group, slot))
747    }
748
749    pub(crate) fn from_backend_buffer<T: TensorScalar, R: TensorRank>(
750        shape: R::Shape,
751        buffer: crate::StorageBuffer<T>,
752    ) -> Result<(Self, DescriptorSlot), GroupError> {
753        let owner = super::root::import_backend_buffer(buffer).map_err(|error| {
754            GroupError::InvalidDescriptor {
755                message: error.to_string(),
756            }
757        })?;
758        let span = owner.root_span();
759        let mut group = Self::new();
760        let allocation = group.insert_owner(owner)?;
761        let layout =
762            TensorLayout::<R>::compact(shape).map_err(|error| GroupError::InvalidDescriptor {
763                message: error.to_string(),
764            })?;
765        let input = DescriptorInput::new(
766            ByteRange::new(0, span.byte_len()),
767            R::shape_from_vec(layout.shape().iter().copied().collect()).map_err(|error| {
768                GroupError::InvalidDescriptor {
769                    message: error.to_string(),
770                }
771            })?,
772            R::strides_from_vec(layout.strides().iter().copied().collect()).map_err(|error| {
773                GroupError::InvalidDescriptor {
774                    message: error.to_string(),
775                }
776            })?,
777            layout.offset(),
778            true,
779        );
780        let slot = group.insert_descriptor::<T, R>(allocation, input)?;
781        Ok((group, slot))
782    }
783
784    pub(crate) fn from_backend_root<T: Send + Sync + 'static>(
785        buffer: crate::StorageBuffer<T>,
786    ) -> Result<Self, GroupError> {
787        let owner = super::root::import_backend_buffer(buffer).map_err(|error| {
788            GroupError::InvalidDescriptor {
789                message: error.to_string(),
790            }
791        })?;
792        let mut group = Self::new();
793        group.insert_owner(owner)?;
794        Ok(group)
795    }
796
797    pub(crate) fn insert_owner(
798        &mut self,
799        owner: OwnedStorage,
800    ) -> Result<AllocationSlot, GroupError> {
801        let index = self.allocations.len();
802        let slot = u32::try_from(index).map_err(|_| GroupError::IndexOverflow)?;
803        self.allocations.push(Some(owner));
804        Ok(AllocationSlot(slot))
805    }
806
807    pub(crate) fn insert_descriptor<T: TensorScalar, R: TensorRank>(
808        &mut self,
809        allocation: AllocationSlot,
810        input: DescriptorInput<R>,
811    ) -> Result<DescriptorSlot, GroupError> {
812        let allocation_index = allocation.index();
813        let owner = self
814            .allocations
815            .get(allocation_index)
816            .ok_or(GroupError::AllocationSlotOutOfBounds {
817                slot: allocation_index,
818            })?
819            .as_ref()
820            .ok_or(GroupError::AllocationSlotVacant {
821                slot: allocation_index,
822            })?;
823
824        let root = owner.as_ref().root_identity();
825        let span = root.bind_relative_range(input.relative).map_err(|error| {
826            GroupError::InvalidDescriptor {
827                message: error.to_string(),
828            }
829        })?;
830        let element_size = size_of::<T>();
831        if element_size == 0 || !span.byte_len().is_multiple_of(element_size) {
832            return Err(GroupError::InvalidDescriptor {
833                message: format!(
834                    "byte span {} is not divisible by element size {}",
835                    span.byte_len(),
836                    element_size
837                ),
838            });
839        }
840        if !span
841            .guaranteed_alignment()
842            .get()
843            .is_multiple_of(align_of::<T>())
844        {
845            return Err(GroupError::InvalidDescriptor {
846                message: format!(
847                    "span alignment {} is insufficient for {}-byte alignment",
848                    span.guaranteed_alignment().get(),
849                    align_of::<T>()
850                ),
851            });
852        }
853
854        let shape = R::shape_into_vec(input.shape);
855        let strides = R::strides_into_vec(input.strides);
856        let layout = TensorLayout::<DynRank>::from_parts(
857            shape,
858            strides,
859            input.offset,
860            span.byte_len() / element_size,
861        )
862        .map_err(|error| GroupError::InvalidDescriptor {
863            message: error.to_string(),
864        })?;
865        let element_count = logical_element_count(layout.shape())?;
866        let write_injective = if input.require_injective {
867            layout.validate_mutable_no_overlap().map_err(|error| {
868                GroupError::InvalidDescriptor {
869                    message: error.to_string(),
870                }
871            })?;
872            true
873        } else {
874            false
875        };
876        let envelope = reachable_envelope(&span, &layout, element_size)?;
877        let (checked, _) = validate_descriptor::<T, DynRank>(
878            &root,
879            span,
880            layout.shape().iter().copied().collect(),
881            layout.strides().iter().copied().collect(),
882            layout.offset(),
883            false,
884        )
885        .map_err(|error| GroupError::InvalidDescriptor {
886            message: error.to_string(),
887        })?;
888        let record = DescriptorRecord {
889            allocation,
890            root,
891            span,
892            layout,
893            dtype: T::dtype(),
894            element_size,
895            element_count,
896            provider: owner.as_ref().provider_kind(),
897            placement: Placement::default(),
898            envelope,
899            write_injective,
900            checked,
901        };
902
903        let descriptor_index = self.descriptors.len();
904        let slot = u32::try_from(descriptor_index).map_err(|_| GroupError::IndexOverflow)?;
905        self.descriptors.push(Some(record));
906        Ok(DescriptorSlot(slot))
907    }
908
909    /// Replace one uniquely-owned descriptor's logical layout while retaining
910    /// its scalar representation and allocation root.
911    // INVARIANT: the dynamic dtype dispatch below selects the sealed scalar
912    // pair before the existing descriptor validator runs.
913    #[allow(clippy::result_large_err)]
914    pub(crate) fn update_descriptor_layout(
915        mut self,
916        slot: DescriptorSlot,
917        shape: Vec<usize>,
918        strides: Vec<isize>,
919        offset: isize,
920    ) -> Result<Self, (Self, GroupError)> {
921        let dtype = match self.resolve_descriptor(slot) {
922            Ok((_, descriptor)) => descriptor.dtype,
923            Err(error) => return Err((self, error)),
924        };
925        if let Some(Some(descriptor)) = self.descriptors.get_mut(slot.index()) {
926            // A metadata-only read view may be non-injective (for example a
927            // broadcast). Mutable access revalidates injectivity when asked.
928            descriptor.write_injective = false;
929        }
930        match dtype {
931            DType::F32 => self.reinterpret_descriptor::<f32, f32>(slot, shape, strides, offset),
932            DType::F64 => self.reinterpret_descriptor::<f64, f64>(slot, shape, strides, offset),
933            DType::I32 => self.reinterpret_descriptor::<i32, i32>(slot, shape, strides, offset),
934            DType::I64 => self.reinterpret_descriptor::<i64, i64>(slot, shape, strides, offset),
935            DType::Bool => self.reinterpret_descriptor::<bool, bool>(slot, shape, strides, offset),
936            DType::C32 => self
937                .reinterpret_descriptor::<num_complex::Complex32, num_complex::Complex32>(
938                    slot, shape, strides, offset,
939                ),
940            DType::C64 => self
941                .reinterpret_descriptor::<num_complex::Complex64, num_complex::Complex64>(
942                    slot, shape, strides, offset,
943                ),
944        }
945    }
946
947    /// Replace one uniquely-owned descriptor with a sealed representation
948    /// reinterpretation while retaining the same allocation root.
949    ///
950    /// The group is returned unchanged with the typed error when validation
951    /// fails, so consuming owner callers can recover the original tensor.
952    // INVARIANT: returning the unchanged move-only group with validation
953    // failure is required so consuming callers can recover its owner.
954    #[allow(clippy::result_large_err)]
955    pub(crate) fn reinterpret_descriptor<T: TensorScalar, U: TensorScalar>(
956        mut self,
957        slot: DescriptorSlot,
958        shape: Vec<usize>,
959        strides: Vec<isize>,
960        offset: isize,
961    ) -> Result<Self, (Self, GroupError)> {
962        let (descriptor_index, descriptor) = match self.resolve_descriptor(slot) {
963            Ok((index, descriptor)) => (index, descriptor.clone()),
964            Err(error) => return Err((self, error)),
965        };
966        if descriptor.dtype != T::dtype() {
967            return Err((
968                self,
969                GroupError::DTypeMismatch {
970                    expected: descriptor.dtype,
971                    actual: T::dtype(),
972                },
973            ));
974        }
975        let allocation = descriptor.allocation;
976        let references = self
977            .descriptors
978            .iter()
979            .flatten()
980            .filter(|candidate| candidate.allocation == allocation)
981            .count();
982        if references != 1 {
983            return Err((
984                self,
985                GroupError::AliasedAllocation {
986                    allocation: allocation.index(),
987                },
988            ));
989        }
990
991        let root = descriptor.root;
992        let span = descriptor.span;
993        let target_element_size = std::mem::size_of::<U>();
994        if target_element_size == 0 || !span.byte_len().is_multiple_of(target_element_size) {
995            return Err((
996                self,
997                GroupError::InvalidDescriptor {
998                    message: format!(
999                        "byte span {} is not divisible by target element size {}",
1000                        span.byte_len(),
1001                        target_element_size
1002                    ),
1003                },
1004            ));
1005        }
1006        let layout = match TensorLayout::<DynRank>::from_parts(
1007            shape.into(),
1008            strides.into(),
1009            offset,
1010            span.byte_len() / target_element_size,
1011        ) {
1012            Ok(layout) => layout,
1013            Err(error) => {
1014                return Err((
1015                    self,
1016                    GroupError::InvalidDescriptor {
1017                        message: error.to_string(),
1018                    },
1019                ))
1020            }
1021        };
1022        let write_injective = descriptor.write_injective;
1023        let envelope = match reachable_envelope(&span, &layout, target_element_size) {
1024            Ok(envelope) => envelope,
1025            Err(error) => return Err((self, error)),
1026        };
1027        let (checked, _) = match validate_descriptor::<U, DynRank>(
1028            &root,
1029            span,
1030            layout.shape().iter().copied().collect(),
1031            layout.strides().iter().copied().collect(),
1032            layout.offset(),
1033            write_injective,
1034        ) {
1035            Ok(value) => value,
1036            Err(error) => {
1037                return Err((
1038                    self,
1039                    GroupError::InvalidDescriptor {
1040                        message: error.to_string(),
1041                    },
1042                ))
1043            }
1044        };
1045        let element_count = match logical_element_count(layout.shape()) {
1046            Ok(count) => count,
1047            Err(error) => return Err((self, error)),
1048        };
1049        self.descriptors[descriptor_index] = Some(DescriptorRecord {
1050            allocation,
1051            root,
1052            span,
1053            layout,
1054            dtype: U::dtype(),
1055            element_size: target_element_size,
1056            element_count,
1057            provider: descriptor.provider,
1058            placement: descriptor.placement.clone(),
1059            envelope,
1060            write_injective,
1061            checked,
1062        });
1063        Ok(self)
1064    }
1065
1066    pub(crate) fn view<T: TensorScalar, R: TensorRank>(
1067        &self,
1068        slot: DescriptorSlot,
1069    ) -> Result<GroupReadView<'_, T, R>, GroupError> {
1070        let (descriptor_index, descriptor) = self.resolve_descriptor(slot)?;
1071        check_typed::<T, R>(descriptor)?;
1072        let owner = self
1073            .allocations
1074            .get(descriptor.allocation.index())
1075            .ok_or(GroupError::AllocationSlotOutOfBounds {
1076                slot: descriptor.allocation.index(),
1077            })?
1078            .as_ref()
1079            .ok_or(GroupError::AllocationSlotVacant {
1080                slot: descriptor.allocation.index(),
1081            })?;
1082        let _ = descriptor_index;
1083        Ok(GroupReadView {
1084            owner: NonNull::from(owner),
1085            descriptor: descriptor.clone(),
1086            _borrow: PhantomData,
1087        })
1088    }
1089
1090    pub(crate) fn view_raw<T: 'static, R: TensorRank>(
1091        &self,
1092        slot: DescriptorSlot,
1093    ) -> Result<GroupReadView<'_, T, R>, GroupError> {
1094        let (_, descriptor) = self.resolve_descriptor(slot)?;
1095        let owner = self
1096            .allocations
1097            .get(descriptor.allocation.index())
1098            .ok_or(GroupError::AllocationSlotOutOfBounds {
1099                slot: descriptor.allocation.index(),
1100            })?
1101            .as_ref()
1102            .ok_or(GroupError::AllocationSlotVacant {
1103                slot: descriptor.allocation.index(),
1104            })?;
1105        Ok(GroupReadView {
1106            owner: NonNull::from(owner),
1107            descriptor: descriptor.clone(),
1108            _borrow: PhantomData,
1109        })
1110    }
1111
1112    pub(crate) fn prepare_device_read_for_layout<T: TensorScalar, R: TensorRank>(
1113        &self,
1114        slot: DescriptorSlot,
1115        layout: &TensorLayout<R>,
1116    ) -> Result<Box<dyn crate::PreparedDeviceAccess + '_>, AccessError> {
1117        self.view_raw::<T, R>(slot)
1118            .map_err(|error| AccessError::InvalidLayout {
1119                message: error.to_string(),
1120            })?
1121            .prepare_device_read_for_layout(layout)
1122    }
1123
1124    pub(crate) fn allocation_index(&self, slot: DescriptorSlot) -> Result<usize, GroupError> {
1125        Ok(self.resolve_descriptor(slot)?.1.allocation.index())
1126    }
1127
1128    pub(crate) fn host_buffer_at<T: 'static>(
1129        &self,
1130        allocation_index: usize,
1131    ) -> Option<&crate::StorageBuffer<T>> {
1132        self.allocations
1133            .get(allocation_index)?
1134            .as_ref()?
1135            .host_buffer::<T>()
1136    }
1137
1138    pub(crate) fn host_root_metadata<T: 'static>(
1139        &self,
1140        slot: DescriptorSlot,
1141    ) -> Option<(usize, usize)> {
1142        let (_, descriptor) = self.resolve_descriptor(slot).ok()?;
1143        let owner = self
1144            .allocations
1145            .get(descriptor.allocation.index())?
1146            .as_ref()?;
1147        if descriptor.span != owner.root_span() {
1148            return None;
1149        }
1150        let crate::StorageBuffer::Host(data) = owner.host_buffer::<T>()? else {
1151            return None;
1152        };
1153        let pointer = data.as_ptr() as usize;
1154        let byte_len = data.len().checked_mul(size_of::<T>())?;
1155        Some((pointer, byte_len))
1156    }
1157
1158    pub(crate) fn backend_buffer<T: 'static>(
1159        &self,
1160        slot: DescriptorSlot,
1161    ) -> Option<&crate::StorageBuffer<T>> {
1162        let (_, descriptor) = self.resolve_descriptor(slot).ok()?;
1163        self.allocations
1164            .get(descriptor.allocation.index())?
1165            .as_ref()?
1166            .backend_buffer::<T>()
1167    }
1168
1169    pub(crate) fn descriptor_len(&self, slot: DescriptorSlot) -> Option<usize> {
1170        self.resolve_descriptor(slot)
1171            .ok()
1172            .map(|(_, descriptor)| descriptor.element_count)
1173    }
1174
1175    pub(crate) fn backend_identity(
1176        &self,
1177        slot: DescriptorSlot,
1178    ) -> Option<(crate::AllocationDomainId, crate::AllocationId)> {
1179        let (_, descriptor) = self.resolve_descriptor(slot).ok()?;
1180        if descriptor.provider == ProviderKind::Cpu {
1181            return None;
1182        }
1183        let owner = self
1184            .allocations
1185            .get(descriptor.allocation.index())?
1186            .as_ref()?;
1187        let key = owner.root_identity().extent().key();
1188        Some((key.domain(), key.local()))
1189    }
1190
1191    pub(crate) fn provider_kind(&self, slot: DescriptorSlot) -> Option<ProviderKind> {
1192        self.resolve_descriptor(slot)
1193            .ok()
1194            .map(|(_, descriptor)| descriptor.provider)
1195    }
1196
1197    pub(crate) fn backend_root_buffer<T: 'static>(&self) -> Option<&crate::StorageBuffer<T>> {
1198        self.allocations
1199            .first()
1200            .and_then(|owner| owner.as_ref())
1201            .and_then(|owner| owner.backend_buffer::<T>())
1202    }
1203
1204    pub(crate) fn backend_buffer_mut<T: 'static>(
1205        &mut self,
1206        slot: DescriptorSlot,
1207    ) -> Option<&mut crate::StorageBuffer<T>> {
1208        let allocation = self.resolve_descriptor(slot).ok()?.1.allocation;
1209        self.allocations
1210            .get_mut(allocation.index())?
1211            .as_mut()?
1212            .backend_buffer_mut::<T>()
1213    }
1214
1215    pub(crate) fn backend_root_buffer_mut<T: 'static>(
1216        &mut self,
1217    ) -> Option<&mut crate::StorageBuffer<T>> {
1218        self.allocations
1219            .first_mut()
1220            .and_then(|owner| owner.as_mut())
1221            .and_then(|owner| owner.backend_buffer_mut::<T>())
1222    }
1223
1224    pub(crate) fn view_mut_raw<T: 'static, R: TensorRank>(
1225        &mut self,
1226        slot: DescriptorSlot,
1227    ) -> Result<GroupWriteView<'_, T, R>, GroupError> {
1228        let (_, descriptor) = self.resolve_descriptor(slot)?;
1229        let descriptor = descriptor.clone();
1230        let owner = self
1231            .allocations
1232            .get_mut(descriptor.allocation.index())
1233            .ok_or(GroupError::AllocationSlotOutOfBounds {
1234                slot: descriptor.allocation.index(),
1235            })?
1236            .as_mut()
1237            .ok_or(GroupError::AllocationSlotVacant {
1238                slot: descriptor.allocation.index(),
1239            })?;
1240        Ok(GroupWriteView {
1241            owner: NonNull::from(owner),
1242            descriptor: descriptor.clone(),
1243            _borrow: PhantomData,
1244        })
1245    }
1246
1247    pub(crate) fn prepare_device_write_for_layout<T: TensorScalar, R: TensorRank>(
1248        &mut self,
1249        slot: DescriptorSlot,
1250        layout: &TensorLayout<R>,
1251    ) -> Result<Box<dyn crate::PreparedDeviceAccess + '_>, AccessError> {
1252        self.view_mut_raw::<T, R>(slot)
1253            .map_err(|error| AccessError::InvalidLayout {
1254                message: error.to_string(),
1255            })?
1256            .prepare_device_write_for_layout(layout)
1257    }
1258
1259    pub(crate) fn view_mut<T: TensorScalar, R: TensorRank>(
1260        &mut self,
1261        slot: DescriptorSlot,
1262    ) -> Result<GroupWriteView<'_, T, R>, GroupError> {
1263        let (descriptor_index, descriptor) = self.resolve_descriptor(slot)?;
1264        let mut descriptor = descriptor.clone();
1265        check_typed::<T, R>(&descriptor)?;
1266        if !descriptor.write_injective {
1267            descriptor
1268                .layout
1269                .validate_mutable_no_overlap()
1270                .map_err(|error| GroupError::InvalidDescriptor {
1271                    message: error.to_string(),
1272                })?;
1273            if let Some(Some(retained)) = self.descriptors.get_mut(descriptor_index) {
1274                retained.write_injective = true;
1275            }
1276            descriptor.write_injective = true;
1277        }
1278        let owner = self
1279            .allocations
1280            .get_mut(descriptor.allocation.index())
1281            .ok_or(GroupError::AllocationSlotOutOfBounds {
1282                slot: descriptor.allocation.index(),
1283            })?
1284            .as_mut()
1285            .ok_or(GroupError::AllocationSlotVacant {
1286                slot: descriptor.allocation.index(),
1287            })?;
1288        Ok(GroupWriteView {
1289            owner: NonNull::from(owner),
1290            descriptor,
1291            _borrow: PhantomData,
1292        })
1293    }
1294
1295    pub(crate) fn split_mut<T: TensorScalar, R: TensorRank>(
1296        &mut self,
1297        slots: &[DescriptorSlot],
1298    ) -> Result<Vec<GroupWriteView<'_, T, R>>, DisjointViewError> {
1299        let mut selected = Vec::with_capacity(slots.len());
1300        for &slot in slots {
1301            let (descriptor_index, descriptor) = self.resolve_descriptor(slot)?;
1302            if selected
1303                .iter()
1304                .any(|(seen, _, _): &(DescriptorSlot, usize, DescriptorRecord)| *seen == slot)
1305            {
1306                return Err(DisjointViewError::DuplicateSlot { slot: slot.index() });
1307            }
1308            check_typed::<T, R>(descriptor)?;
1309            if !descriptor.write_injective {
1310                descriptor
1311                    .layout
1312                    .validate_mutable_no_overlap()
1313                    .map_err(|_| DisjointViewError::NonInjective { slot: slot.index() })?;
1314            }
1315            selected.push((slot, descriptor_index, descriptor.clone()));
1316        }
1317
1318        for left in 0..selected.len() {
1319            for right in (left + 1)..selected.len() {
1320                let first = &selected[left].2;
1321                let second = &selected[right].2;
1322                if first.root.root_resource() != second.root.root_resource() {
1323                    continue;
1324                }
1325                match (first.envelope, second.envelope) {
1326                    (None, _) | (_, None) => {}
1327                    (Some(first), Some(second)) => {
1328                        if first
1329                            .overlaps(second)
1330                            .map_err(|_| DisjointViewError::NotProvablyDisjoint)?
1331                        {
1332                            return Err(DisjointViewError::PairwiseOverlap);
1333                        }
1334                    }
1335                }
1336            }
1337        }
1338
1339        for (_, descriptor_index, _) in &selected {
1340            if let Some(Some(descriptor)) = self.descriptors.get_mut(*descriptor_index) {
1341                descriptor.write_injective = true;
1342            }
1343        }
1344
1345        let mut children = Vec::with_capacity(selected.len());
1346        for (_, _, mut descriptor) in selected {
1347            descriptor.write_injective = true;
1348            let owner = self
1349                .allocations
1350                .get_mut(descriptor.allocation.index())
1351                .ok_or(GroupError::AllocationSlotOutOfBounds {
1352                    slot: descriptor.allocation.index(),
1353                })?
1354                .as_mut()
1355                .ok_or(GroupError::AllocationSlotVacant {
1356                    slot: descriptor.allocation.index(),
1357                })?;
1358            children.push(GroupWriteView {
1359                owner: NonNull::from(owner),
1360                descriptor,
1361                _borrow: PhantomData,
1362            });
1363        }
1364        Ok(children)
1365    }
1366
1367    pub(crate) fn try_extract(
1368        &mut self,
1369        slot: DescriptorSlot,
1370    ) -> Result<OwnedStorage, ExtractError> {
1371        let (descriptor_index, descriptor) = self.resolve_descriptor(slot)?;
1372        let allocation = descriptor.allocation;
1373        let references = self
1374            .descriptors
1375            .iter()
1376            .flatten()
1377            .filter(|candidate| candidate.allocation == allocation)
1378            .count();
1379        if references != 1 {
1380            return Err(ExtractError::AliasedAllocation {
1381                allocation: allocation.index(),
1382            });
1383        }
1384        let _ = self.descriptors[descriptor_index].take();
1385        self.allocations
1386            .get_mut(allocation.index())
1387            .ok_or(GroupError::AllocationSlotOutOfBounds {
1388                slot: allocation.index(),
1389            })?
1390            .take()
1391            .ok_or(GroupError::AllocationSlotVacant {
1392                slot: allocation.index(),
1393            })
1394            .map_err(ExtractError::from)
1395    }
1396
1397    /// Extract one uniquely-owned descriptor while retaining all other group
1398    /// descriptors in place.
1399    ///
1400    /// This is structural: aliased allocations return a typed error and the
1401    /// group remains unchanged.
1402    ///
1403    /// # Errors
1404    ///
1405    /// Returns [`GroupError::AliasedAllocation`] for an aliased allocation or
1406    /// [`GroupError::InvalidDescriptor`] for an invalid slot; every extraction
1407    /// failure leaves the group unchanged.
1408    #[allow(clippy::result_large_err)]
1409    pub fn take_tensor(&mut self, slot: DescriptorSlot) -> Result<crate::Tensor, GroupError> {
1410        let (_, descriptor) = self.resolve_descriptor(slot)?;
1411        let descriptor = descriptor.clone();
1412        let dtype = descriptor.dtype;
1413        let layout = descriptor.layout.clone();
1414        let placement = descriptor.placement.clone();
1415        let owner = self
1416            .try_extract(slot)
1417            .map_err(|error| GroupError::InvalidDescriptor {
1418                message: error.to_string(),
1419            })?;
1420        let mut extracted = Self::new();
1421        extracted.allocations.push(Some(owner));
1422        let mut descriptor = descriptor.clone();
1423        descriptor.allocation = AllocationSlot(0);
1424        extracted.descriptors.push(Some(descriptor));
1425        Ok(tensor_from_group(
1426            extracted,
1427            DescriptorSlot(0),
1428            0,
1429            dtype,
1430            layout,
1431            placement,
1432        ))
1433    }
1434
1435    // INVARIANT: extraction failure returns the unchanged group because the
1436    // caller must retain ownership when a descriptor cannot be detached.
1437    #[allow(clippy::result_large_err)]
1438    pub(crate) fn into_owner(
1439        mut self,
1440        slot: DescriptorSlot,
1441    ) -> Result<OwnedStorage, (Self, ExtractError)> {
1442        let result = self.try_extract(slot);
1443        match result {
1444            Ok(owner) => Ok(owner),
1445            Err(error) => Err((self, error)),
1446        }
1447    }
1448
1449    // INVARIANT: returning the unchanged move-only group is the extraction
1450    // failure carrier required by the ownership contract.
1451    /// Consume the group and extract one descriptor as a standalone tensor.
1452    ///
1453    /// Extraction is structural: it succeeds only when no other descriptor
1454    /// aliases the selected physical allocation. Failure returns the exact
1455    /// unchanged group and typed error without copying.
1456    ///
1457    /// # Errors
1458    ///
1459    /// Returns [`GroupError::AliasedAllocation`] when another descriptor
1460    /// references the allocation, or [`GroupError::InvalidDescriptor`] for an
1461    /// invalid slot. Each extraction failure returns the unchanged group.
1462    #[allow(clippy::result_large_err)]
1463    pub fn into_tensor(self, slot: DescriptorSlot) -> Result<crate::Tensor, (Self, GroupError)> {
1464        let (_, descriptor) = match self.resolve_descriptor(slot) {
1465            Ok(value) => value,
1466            Err(error) => return Err((self, error)),
1467        };
1468        let allocation = descriptor.allocation;
1469        let references = self
1470            .descriptors
1471            .iter()
1472            .flatten()
1473            .filter(|candidate| candidate.allocation == allocation)
1474            .count();
1475        if references != 1 {
1476            return Err((
1477                self,
1478                GroupError::AliasedAllocation {
1479                    allocation: allocation.index(),
1480                },
1481            ));
1482        }
1483        let dtype = descriptor.dtype;
1484        let layout = descriptor.layout.clone();
1485        let placement = descriptor.placement.clone();
1486        let (group, slot) = match self.into_single_descriptor(slot) {
1487            Ok(value) => value,
1488            Err((group, error)) => return Err((group, error)),
1489        };
1490        Ok(tensor_from_group(group, slot, 0, dtype, layout, placement))
1491    }
1492
1493    // INVARIANT: structural extraction must return the unchanged group on
1494    // every validation failure so no owner is lost.
1495    #[allow(clippy::result_large_err)]
1496    fn into_single_descriptor(
1497        mut self,
1498        slot: DescriptorSlot,
1499    ) -> Result<(Self, DescriptorSlot), (Self, GroupError)> {
1500        let (descriptor_index, descriptor) = match self.resolve_descriptor(slot) {
1501            Ok(value) => value,
1502            Err(error) => return Err((self, error)),
1503        };
1504        let allocation = descriptor.allocation;
1505        let mut descriptor = match self.descriptors[descriptor_index].take() {
1506            Some(descriptor) => descriptor,
1507            None => {
1508                return Err((
1509                    self,
1510                    GroupError::DescriptorSlotVacant { slot: slot.index() },
1511                ))
1512            }
1513        };
1514        let owner = match self.allocations.get_mut(allocation.index()) {
1515            Some(owner) => match owner.take() {
1516                Some(owner) => owner,
1517                None => {
1518                    self.descriptors[descriptor_index] = Some(descriptor);
1519                    return Err((
1520                        self,
1521                        GroupError::AllocationSlotVacant {
1522                            slot: allocation.index(),
1523                        },
1524                    ));
1525                }
1526            },
1527            None => {
1528                self.descriptors[descriptor_index] = Some(descriptor);
1529                return Err((
1530                    self,
1531                    GroupError::AllocationSlotOutOfBounds {
1532                        slot: allocation.index(),
1533                    },
1534                ));
1535            }
1536        };
1537        let mut group = Self::new();
1538        descriptor.allocation = AllocationSlot(0);
1539        group.allocations.push(Some(owner));
1540        group.descriptors.push(Some(descriptor));
1541        Ok((group, DescriptorSlot(0)))
1542    }
1543
1544    pub(crate) fn into_host_vec<T: TensorScalar>(
1545        self,
1546        slot: DescriptorSlot,
1547    ) -> Result<Vec<T>, String> {
1548        let owner = self
1549            .into_owner(slot)
1550            .map_err(|(_, error)| error.to_string())?;
1551        owner
1552            .into_host_vec::<T>()
1553            .map_err(|error| error.to_string())
1554    }
1555
1556    fn resolve_descriptor(
1557        &self,
1558        slot: DescriptorSlot,
1559    ) -> Result<(usize, &DescriptorRecord), GroupError> {
1560        let index = slot.index();
1561        let descriptor = self
1562            .descriptors
1563            .get(index)
1564            .ok_or(GroupError::DescriptorSlotOutOfBounds { slot: index })?
1565            .as_ref()
1566            .ok_or(GroupError::DescriptorSlotVacant { slot: index })?;
1567        Ok((index, descriptor))
1568    }
1569
1570    #[cfg(test)]
1571    pub(crate) fn test_vacate_allocation(&mut self, slot: AllocationSlot) {
1572        if let Some(entry) = self.allocations.get_mut(slot.index()) {
1573            *entry = None;
1574        }
1575    }
1576}
1577
1578fn check_typed<T: TensorScalar, R: TensorRank>(
1579    descriptor: &DescriptorRecord,
1580) -> Result<(), GroupError> {
1581    if descriptor.dtype != T::dtype() {
1582        return Err(GroupError::DTypeMismatch {
1583            expected: descriptor.dtype,
1584            actual: T::dtype(),
1585        });
1586    }
1587    if let Some(expected) = R::RANK {
1588        let actual = descriptor.layout.shape().len();
1589        if expected != actual {
1590            return Err(GroupError::RankMismatch { expected, actual });
1591        }
1592    }
1593    Ok(())
1594}
1595
1596fn logical_element_count(shape: &[usize]) -> Result<usize, GroupError> {
1597    shape.iter().try_fold(1usize, |count, &extent| {
1598        count
1599            .checked_mul(extent)
1600            .ok_or_else(|| GroupError::InvalidDescriptor {
1601                message: "logical element count overflows".to_owned(),
1602            })
1603    })
1604}
1605
1606fn reachable_envelope(
1607    span: &RootBoundSpan,
1608    layout: &TensorLayout<DynRank>,
1609    element_size: usize,
1610) -> Result<Option<ByteRange>, GroupError> {
1611    if layout.shape().contains(&0) {
1612        return Ok(None);
1613    }
1614    let mut minimum = layout.offset() as i128;
1615    let mut maximum = minimum;
1616    for (&extent, &stride) in layout.shape().iter().zip(layout.strides()) {
1617        let steps = i128::try_from(extent - 1).map_err(|_| GroupError::InvalidDescriptor {
1618            message: "layout extent does not fit i128".to_owned(),
1619        })?;
1620        let contribution =
1621            (stride as i128)
1622                .checked_mul(steps)
1623                .ok_or_else(|| GroupError::InvalidDescriptor {
1624                    message: "reachable layout arithmetic overflows".to_owned(),
1625                })?;
1626        if contribution < 0 {
1627            minimum =
1628                minimum
1629                    .checked_add(contribution)
1630                    .ok_or_else(|| GroupError::InvalidDescriptor {
1631                        message: "reachable layout minimum overflows".to_owned(),
1632                    })?;
1633        } else {
1634            maximum =
1635                maximum
1636                    .checked_add(contribution)
1637                    .ok_or_else(|| GroupError::InvalidDescriptor {
1638                        message: "reachable layout maximum overflows".to_owned(),
1639                    })?;
1640        }
1641    }
1642    let minimum = usize::try_from(minimum).map_err(|_| GroupError::InvalidDescriptor {
1643        message: "reachable layout minimum is negative".to_owned(),
1644    })?;
1645    let maximum = usize::try_from(maximum).map_err(|_| GroupError::InvalidDescriptor {
1646        message: "reachable layout maximum is negative or too large".to_owned(),
1647    })?;
1648    let byte_offset = span
1649        .byte_offset()
1650        .checked_add(minimum.checked_mul(element_size).ok_or_else(|| {
1651            GroupError::InvalidDescriptor {
1652                message: "reachable byte offset overflows".to_owned(),
1653            }
1654        })?)
1655        .ok_or_else(|| GroupError::InvalidDescriptor {
1656            message: "reachable byte offset overflows".to_owned(),
1657        })?;
1658    let byte_len = maximum
1659        .checked_sub(minimum)
1660        .and_then(|length| length.checked_add(1))
1661        .and_then(|length| length.checked_mul(element_size))
1662        .ok_or_else(|| GroupError::InvalidDescriptor {
1663            message: "reachable byte length overflows".to_owned(),
1664        })?;
1665    let range = ByteRange::new(byte_offset, byte_len);
1666    range
1667        .checked_end()
1668        .map_err(|error| GroupError::InvalidDescriptor {
1669            message: error.to_string(),
1670        })?;
1671    Ok(Some(range))
1672}
1673
1674#[cfg(test)]
1675pub(crate) fn test_logical_element_count(shape: &[usize]) -> Result<usize, GroupError> {
1676    logical_element_count(shape)
1677}
1678
1679#[cfg(test)]
1680pub(crate) fn test_reachable_envelope(
1681    span: RootBoundSpan,
1682    layout: TensorLayout<DynRank>,
1683    element_size: usize,
1684) -> Result<Option<ByteRange>, GroupError> {
1685    reachable_envelope(&span, &layout, element_size)
1686}