fuel-ir 0.20.60

// SPDX-License-Identifier: MIT OR Apache-2.0
//! Tensor memory layout: shape, strides, and start offset.
//!
//! A [`Layout`] describes how a tensor's elements are stored in memory.
//! Contiguous layouts use row-major strides; non-contiguous layouts arise from
//! operations like [`transpose`](Layout::transpose) or [`narrow`](Layout::narrow).
use crate::{DimVec, Error, Result, Shape, StrideVec};
use smallvec::smallvec;

/// Describes the memory layout of a tensor: shape, strides, and start offset.
///
/// Strides are in units of elements (not bytes) and are **signed** —
/// metadata-only view ops (notably `Op::Flip`) reverse iteration along
/// a dim by negating that dim's stride and adjusting `start_offset`,
/// so a stride can be negative even though sizes can't. A contiguous
/// (row-major) tensor of shape `[a, b, c]` has strides `[b*c, c, 1]`
/// (all positive); a Flip-on-dim-0 of that has stride `[-b*c, c, 1]`
/// with `start_offset` set to the byte offset of the (originally)
/// last element along dim 0.
///
/// `start_offset` is the byte offset of the layout's *first* element
/// in iteration order — by construction always non-negative, even for
/// negative-stride layouts.
#[derive(Debug, PartialEq, Eq, Clone)]
pub struct Layout {
    shape: Shape,
    // Signed: see struct doc comment.
    stride: StrideVec,
    start_offset: usize,
}

impl Layout {
    /// Creates a layout with explicit shape, strides, and start offset.
    pub fn new(shape: Shape, stride: StrideVec, start_offset: usize) -> Self {
        Self {
            shape,
            stride,
            start_offset,
        }
    }

    /// Creates a contiguous (row-major) layout with the given start offset.
    pub fn contiguous_with_offset<S: Into<Shape>>(shape: S, start_offset: usize) -> Self {
        let shape = shape.into();
        let stride = shape.stride_contiguous();
        Self {
            shape,
            stride,
            start_offset,
        }
    }

    /// Creates a contiguous (row-major) layout with zero start offset.
    pub fn contiguous<S: Into<Shape>>(shape: S) -> Self {
        Self::contiguous_with_offset(shape, 0)
    }

    /// Returns the dimension sizes as a slice.
    pub fn dims(&self) -> &[usize] {
        self.shape.dims()
    }

    /// Returns the size of a specific dimension.
    pub fn dim<D: crate::shape::Dim>(&self, dim: D) -> Result<usize> {
        let dim = dim.to_index(&self.shape, "dim")?;
        Ok(self.dims()[dim])
    }

    /// Returns a reference to the [`Shape`].
    pub fn shape(&self) -> &Shape {
        &self.shape
    }

    /// Whether the layout's shape has any bounded-symbolic (`Range`) axis.
    /// Phase D step 1b.
    pub fn has_dynamic(&self) -> bool {
        self.shape.has_dynamic()
    }

    /// Resolve the layout's shape through `env` (every dynamic axis →
    /// concrete), keeping **strides and offset unchanged**: the physical walk
    /// of a fixed-capacity buffer is the same; only the logical live extent
    /// becomes concrete. Errors if a dynamic axis's symbol is unbound.
    /// Phase D step 1b.
    pub fn resolve(&self, env: &crate::SymEnv) -> Result<Layout> {
        Ok(Layout {
            shape: self.shape.resolve(env)?,
            stride: self.stride.clone(),
            start_offset: self.start_offset,
        })
    }

    /// Returns the strides as a slice (in elements, not bytes).
    /// Signed — see struct-level docs.
    pub fn stride(&self) -> &[isize] {
        &self.stride
    }

    /// Returns strides cast to `usize`, debug-asserting non-negative.
    /// Convenience for kernels that haven't been updated to handle
    /// signed strides yet (most of them, since auto-Contiguize
    /// guarantees positive strides on kernel input today). One stack
    /// allocation via [`DimVec`]; effectively zero cost for rank ≤ 6.
    ///
    /// **Don't use this** if your kernel needs to support negative
    /// strides (e.g. is the consumer of an `Op::Flip` view): iterate
    /// via [`crate::StridedIndex`] which handles the sign correctly.
    pub fn stride_unsigned(&self) -> DimVec {
        self.stride
            .iter()
            .map(|&s| {
                debug_assert!(
                    s >= 0,
                    "stride_unsigned() called on a layout with negative stride; \
                use StridedIndex or operate on signed strides"
                );
                s as usize
            })
            .collect()
    }

    /// Returns the start offset into the underlying storage.
    pub fn start_offset(&self) -> usize {
        self.start_offset
    }

    /// Returns `(start, end)` offsets if the layout is contiguous, `None` otherwise.
    pub fn contiguous_offsets(&self) -> Option<(usize, usize)> {
        if self.is_contiguous() {
            let start_o = self.start_offset;
            Some((start_o, start_o + self.shape.elem_count()))
        } else {
            None
        }
    }

    /// Returns `true` if the layout is C-contiguous (row-major).
    ///
    /// This does not require the start offset to be 0.
    pub fn is_contiguous(&self) -> bool {
        self.shape.is_contiguous(&self.stride)
    }

    /// Returns `true` if the layout is Fortran-contiguous (column-major).
    pub fn is_fortran_contiguous(&self) -> bool {
        self.shape.is_fortran_contiguous(&self.stride)
    }

    /// Returns a narrowed layout selecting `len` elements starting at `start` on dimension `dim`.
    pub fn narrow(&self, dim: usize, start: usize, len: usize) -> Result<Self> {
        let dims = self.shape().dims();
        if dim >= dims.len() {
            Err(Error::DimOutOfRange {
                shape: Box::new(self.shape().clone()),
                dim: dim as i32,
                op: "narrow",
            }
            .bt())?
        }
        if start + len > dims[dim] {
            Err(Error::NarrowInvalidArgs {
                shape: Box::new(self.shape.clone()),
                dim,
                start,
                len,
                msg: "start + len > dim_len",
            }
            .bt())?
        }
        let mut dims = DimVec::from_slice(dims);
        dims[dim] = len;
        // start * stride[dim] is element offset; for negative strides
        // it shifts in the negative element direction, which then adds
        // (potentially as a wrapping subtraction) to start_offset.
        // start_offset stays non-negative because narrow only ever
        // moves the iteration window forward in the original storage's
        // sense — for positive stride, start*stride is non-negative;
        // for negative stride, start*stride is non-positive and we add
        // a non-positive shift, which is fine as long as the original
        // start_offset already accommodates the layout's full footprint.
        let new_start = (self.start_offset as isize + self.stride[dim] * start as isize) as usize;
        Ok(Self {
            shape: Shape::from(dims),
            stride: self.stride.clone(),
            start_offset: new_start,
        })
    }

    /// Returns a layout with two dimensions swapped.
    pub fn transpose(&self, dim1: usize, dim2: usize) -> Result<Self> {
        let rank = self.shape.rank();
        if rank <= dim1 || rank <= dim2 {
            Err(Error::UnexpectedNumberOfDims {
                expected: usize::max(dim1, dim2),
                got: rank,
                shape: Box::new(self.shape().clone()),
            }
            .bt())?
        }
        let mut stride = StrideVec::from_slice(self.stride());
        let mut dims = DimVec::from_slice(self.shape().dims());
        dims.swap(dim1, dim2);
        stride.swap(dim1, dim2);
        Ok(Self {
            shape: Shape::from(dims),
            stride,
            start_offset: self.start_offset,
        })
    }

    /// Returns a layout with dimensions reordered according to `idxs`.
    pub fn permute(&self, idxs: &[usize]) -> Result<Self> {
        let is_permutation =
            idxs.len() == self.shape.rank() && (0..idxs.len()).all(|i| idxs.contains(&i));
        if !is_permutation {
            crate::bail!(
                "dimension mismatch in permute, tensor {:?}, dims: {:?}",
                self.dims(),
                idxs
            )
        }
        let stride = self.stride();
        let dims = self.shape().dims();
        let mut perm_stride = StrideVec::from_slice(stride);
        let mut perm_dims = DimVec::from_slice(dims);
        for (i, &idx) in idxs.iter().enumerate() {
            perm_stride[i] = stride[idx];
            perm_dims[i] = dims[idx];
        }
        Ok(Self {
            shape: Shape::from(perm_dims),
            stride: perm_stride,
            start_offset: self.start_offset,
        })
    }

    /// Returns a layout with a size-1 dimension inserted at position `dim`.
    /// `dim` must be in `0..=rank` — `dim == rank` appends to the end.
    /// Pure metadata: the underlying bytes are unchanged.
    ///
    /// Stride for the inserted size-1 axis is set to 0. This stride
    /// value never affects indexing (the axis only ever has coordinate
    /// 0, and `0 * stride == 0`), so any value would be correct.
    /// Choosing 0 matches the convention used by `broadcast_as` for
    /// size-1 axes broadcast to larger sizes — keeps the convention
    /// uniform and `Shape::is_contiguous` stays unaffected (it already
    /// skips dims of size 1).
    pub fn unsqueeze(&self, dim: usize) -> Result<Self> {
        let rank = self.shape().rank();
        if dim > rank {
            return Err(Error::Msg(format!(
                "unsqueeze: dim {dim} out of bounds for rank {rank}",
            ))
            .bt());
        }
        let mut new_dims = DimVec::from_slice(self.shape().dims());
        new_dims.insert(dim, 1);
        let mut new_stride = StrideVec::from_slice(self.stride());
        new_stride.insert(dim, 0);
        Ok(Self {
            shape: Shape::from(new_dims),
            stride: new_stride,
            start_offset: self.start_offset,
        })
    }

    /// Drop a size-1 dimension at position `dim`. Inverse of
    /// [`Self::unsqueeze`]. The pruned dim's stride is discarded —
    /// it's irrelevant for indexing because the dim has size 1.
    /// Errors if `dim` is out of bounds or its size isn't 1.
    pub fn squeeze(&self, dim: usize) -> Result<Self> {
        let rank = self.shape().rank();
        if dim >= rank {
            return Err(
                Error::Msg(format!("squeeze: dim {dim} out of bounds for rank {rank}",)).bt(),
            );
        }
        if self.shape().dims()[dim] != 1 {
            return Err(Error::Msg(format!(
                "squeeze: dim {dim} has size {}, expected 1",
                self.shape().dims()[dim],
            ))
            .bt());
        }
        let mut new_dims = DimVec::from_slice(self.shape().dims());
        new_dims.remove(dim);
        let mut new_stride = StrideVec::from_slice(self.stride());
        new_stride.remove(dim);
        Ok(Self {
            shape: Shape::from(new_dims),
            stride: new_stride,
            start_offset: self.start_offset,
        })
    }

    /// Returns a layout that iterates `dim` in reverse — pure metadata,
    /// zero copy. The stride at `dim` is negated; `start_offset` is
    /// shifted forward by `(size-1) * |stride|` so iteration starts at
    /// what was the last element along `dim`.
    ///
    /// Size-0 and size-1 axes are no-ops (no element to flip; the new
    /// start_offset shift is 0). Successive `flip` calls on the same
    /// dim are involutive: stride re-negates, offset cancels, you get
    /// the original layout back.
    pub fn flip(&self, dim: usize) -> Result<Self> {
        let rank = self.shape().rank();
        if dim >= rank {
            return Err(Error::Msg(format!("flip: dim {dim} out of bounds for rank {rank}",)).bt());
        }
        let size = self.shape().dims()[dim];
        let mut new_stride = StrideVec::from_slice(self.stride());
        let s = new_stride[dim];
        new_stride[dim] = -s;
        // For size <= 1 the iteration window is a single position (or
        // empty), so the offset shift is 0 — flipping degenerate axes
        // is a no-op.
        let shift = if size > 1 { (size as isize - 1) * s } else { 0 };
        // shift may be negative when stride is positive (most common
        // case) — that pushes start_offset forward in storage units.
        // For an already-negative stride we'd shift backwards, which
        // is exactly what undoes a prior flip.
        let new_start = (self.start_offset as isize + shift) as usize;
        Ok(Self {
            shape: self.shape.clone(),
            stride: new_stride,
            start_offset: new_start,
        })
    }

    /// Returns a layout broadcast to the given shape.
    ///
    /// Dimensions of size 1 in `self` are stretched to match the target shape.
    pub fn broadcast_as<S: Into<Shape>>(&self, shape: S) -> Result<Self> {
        let shape = shape.into();
        if shape.rank() < self.shape().rank() {
            return Err(Error::BroadcastIncompatibleShapes {
                src_shape: Box::new(self.shape().clone()),
                dst_shape: Box::new(shape),
            }
            .bt());
        }
        let added_dims = shape.rank() - self.shape().rank();
        let mut stride: StrideVec = smallvec![0_isize; added_dims];
        for (&dst_dim, (&src_dim, &src_stride)) in shape.dims()[added_dims..]
            .iter()
            .zip(self.dims().iter().zip(self.stride()))
        {
            let s = if dst_dim == src_dim {
                src_stride
            } else if src_dim != 1 {
                return Err(Error::BroadcastIncompatibleShapes {
                    src_shape: Box::new(self.shape().clone()),
                    dst_shape: Box::new(shape),
                }
                .bt());
            } else {
                0
            };
            stride.push(s)
        }
        Ok(Self {
            shape,
            stride,
            start_offset: self.start_offset,
        })
    }

    pub fn strided_index(&self) -> crate::StridedIndex<'_> {
        crate::StridedIndex::from_layout(self)
    }

    pub fn strided_blocks(&self) -> crate::StridedBlocks<'_> {
        let mut block_len: usize = 1;
        let mut contiguous_dims = 0; // These are counted from the right.
        for (&stride, &dim) in self.stride().iter().zip(self.dims().iter()).rev() {
            // Contiguous block detection: stride must equal +block_len
            // exactly. Negative strides break the run (a flipped axis
            // can't extend a contiguous block).
            if stride != block_len as isize {
                break;
            }
            block_len *= dim;
            contiguous_dims += 1;
        }
        let index_dims = self.dims().len() - contiguous_dims;
        if index_dims == 0 {
            crate::StridedBlocks::SingleBlock {
                start_offset: self.start_offset,
                len: block_len,
            }
        } else {
            let block_start_index = crate::StridedIndex::new(
                &self.dims()[..index_dims],
                &self.stride[..index_dims],
                self.start_offset,
            );
            crate::StridedBlocks::MultipleBlocks {
                block_start_index,
                block_len,
            }
        }
    }

    // Returns the contiguous offsets with broadcast if applicable.
    pub fn offsets_b(&self) -> Option<ContiguousOffsetsWithBroadcast> {
        let mut left_broadcast: usize = 1;
        let mut right_broadcast: usize = 1;
        let strides = self.stride();
        let dims = self.dims();
        let mut start_cont = 0;
        let mut end_cont = dims.len();
        for (&s, &d) in strides.iter().zip(dims.iter()) {
            if s != 0 {
                break;
            }
            start_cont += 1;
            left_broadcast *= d;
        }
        if start_cont == dims.len() {
            return Some(ContiguousOffsetsWithBroadcast {
                start: self.start_offset,
                len: 1,
                left_broadcast,
                right_broadcast: 1,
            });
        }
        for (&s, &d) in strides.iter().zip(dims.iter()).rev() {
            if s != 0 {
                break;
            }
            end_cont -= 1;
            right_broadcast *= d;
        }
        // Check that the inner dims are contiguous
        let strides = &strides[start_cont..end_cont];
        let dims = &dims[start_cont..end_cont];
        let mut len: usize = 1;
        for (&stride, &dim) in strides.iter().zip(dims.iter()).rev() {
            // Same as strided_blocks: positive contiguous run only.
            if stride != len as isize {
                return None;
            }
            len *= dim;
        }
        Some(ContiguousOffsetsWithBroadcast {
            start: self.start_offset,
            len,
            left_broadcast,
            right_broadcast,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn flip_dim0_of_3x4_negates_stride_and_shifts_offset() {
        let l = Layout::contiguous(Shape::from_dims(&[3, 4]));
        let f = l.flip(0).unwrap();
        assert_eq!(f.shape().dims(), &[3, 4]);
        assert_eq!(f.stride(), &[-4_isize, 1]);
        assert_eq!(f.start_offset(), 8); // (3-1) * 4
        assert!(!f.is_contiguous());
    }

    #[test]
    fn flip_is_involutive() {
        let l = Layout::contiguous(Shape::from_dims(&[3, 4]));
        let ff = l.flip(0).unwrap().flip(0).unwrap();
        assert_eq!(ff.shape().dims(), l.shape().dims());
        assert_eq!(ff.stride(), l.stride());
        assert_eq!(ff.start_offset(), l.start_offset());
    }

    #[test]
    fn flip_size_one_axis_is_noop_offset() {
        let l = Layout::contiguous(Shape::from_dims(&[1, 4]));
        let f = l.flip(0).unwrap();
        assert_eq!(f.start_offset(), 0);
        assert_eq!(f.stride(), &[-4_isize, 1]);
    }

    #[test]
    fn flip_iteration_yields_reversed_offsets() {
        let l = Layout::contiguous(Shape::from_dims(&[4]));
        let f = l.flip(0).unwrap();
        let offsets: Vec<usize> = f.strided_index().collect();
        assert_eq!(offsets, vec![3, 2, 1, 0]);
    }

    #[test]
    fn flip_dim_oob_errors() {
        let l = Layout::contiguous(Shape::from_dims(&[3, 4]));
        assert!(l.flip(2).is_err());
    }
}

/// Offsets for a layout that is contiguous in its inner dimensions but may be broadcast
/// on the left or right outer dimensions.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ContiguousOffsetsWithBroadcast {
    /// Start offset of the contiguous block.
    pub start: usize,
    /// Length of the contiguous block in elements.
    pub len: usize,
    /// Number of times the block is repeated on the left (outer) dimensions.
    pub left_broadcast: usize,
    /// Number of times the block is repeated on the right (inner) dimensions.
    pub right_broadcast: usize,
}