UnmanagedStorage.Slicing.cs

using System;
using System.Collections;
using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
using System.Linq;
using System.Runtime.CompilerServices;
using NumSharp.Backends.Unmanaged;

namespace NumSharp.Backends
{
    public partial class UnmanagedStorage
    {
        #region Slicing

        [MethodImpl(Inline)]
        public UnmanagedStorage GetView(string slicing_notation) => GetView(Slice.ParseSlices(slicing_notation));

        [MethodImpl(Inline)]
        public UnmanagedStorage GetView(params Slice[] slices)
        {
            if (slices == null)
                throw new ArgumentNullException(nameof(slices));
            // deal with ellipsis and newaxis if any before continuing into GetViewInternal
            int ellipsis_count = 0;
            int newaxis_count = 0;
            foreach (var slice in slices)
            {
                if (slice.IsEllipsis)
                    ellipsis_count++;
                if (slice.IsNewAxis)
                    newaxis_count++;
            }

            // Too many indices: count the AXIS-CONSUMING slices (everything except newaxis
            // and ellipsis). NumPy raises IndexError when this exceeds the array's rank — an
            // ellipsis only ABSORBS axes, so it can never rescue an already-over-full index
            // (A[:, :, :] / A[..., :, :, :] on a 2-D array both raise). Previously the extra
            // slices were silently dropped by Shape.Slice (it iterates only over NDim).
            int consuming = slices.Length - newaxis_count - ellipsis_count;
            if (consuming > Shape.NDim)
                throw new IndexError($"too many indices for array: array is {Shape.NDim}-dimensional, but {consuming} were indexed");

            // deal with ellipsis
            if (ellipsis_count > 1)
                throw new ArgumentException("IndexError: an index can only have a single ellipsis ('...')");
            else if (ellipsis_count == 1)
                slices = ExpandEllipsis(slices).ToArray();

            // deal with newaxis
            if (newaxis_count > 0)
            {
                var view = this;
                for (var axis = 0; axis < slices.Length; axis++)
                {
                    var slice = slices[axis];
                    if (slice.IsNewAxis)
                    {
                        slices[axis] = Slice.All;
                        view = view.Alias(view.Shape.ExpandDimension(axis));
                    }
                }

                return view.GetViewInternal(slices);
            }

            // slicing without newaxis
            return GetViewInternal(slices);
        }

        [MethodImpl(Inline)]
        [SuppressMessage("ReSharper", "PossibleInvalidOperationException")]
        private UnmanagedStorage GetViewInternal(params Slice[] slices)
        {
            // NOTE: GetViewInternal can not deal with Slice.Ellipsis or Slice.NewAxis!
            // For contiguous shapes, if all slices are indices, we can use GetData to get
            // a direct memory slice. For non-contiguous shapes (stepped, transposed), we
            // must go through the full slice path to create a proper view.
            if (_shape.IsContiguous)
            {
                var indices = new long[slices.Length];
                for (var i = 0; i < slices.Length; i++)
                {
                    var inputSlice = slices[i];
                    if (!inputSlice.IsIndex)
                    {
                        //incase it is a trailing :, e.g. [2,2, :] in a shape (3,3,5,5) -> (5,5)
                        if (i == slices.Length - 1 && inputSlice == Slice.All)
                        {
                            Array.Resize(ref indices, indices.Length - 1);
                            goto _getdata;
                        }

                        goto _perform_slice;
                    }

                    indices[i] = inputSlice.Start.Value;
                }

                _getdata:
                return GetData(indices);
            }

            //perform a regular slicing
            _perform_slice:

            // In case the slices selected are all ":" - just return an alias
            if (slices.All(s => Equals(Slice.All, s)))
                return Alias();

            // Broadcast arrays: DO NOT materialize! Shape.Slice correctly handles stride=0.
            // The old code materialized broadcast data, losing stride=0 information.
            // NumPy preserves stride=0 when slicing broadcast views.
            // Example: broadcast_to([1,2,3], (4,3)) with strides (0,1)
            //   [:, 0] should give shape (4,) with stride (0,) - NOT materialized stride (3,)
            //
            // NOTE: Cumsum and other axis reductions need this to work correctly with
            // broadcast arrays without requiring memory-expensive .copy() workarounds.

            var slicedShape = _shape.Slice(slices);

            // Handle empty slices (e.g., a[100:200] on a 10-element array)
            // NumPy returns an empty array with shape (0,) or similar
            if (slicedShape.size == 0)
            {
                // Create empty storage with correct dtype and shape. It is still a VIEW of the parent
                // (NumPy: owndata=False, base set — there are no elements, so the fresh zero-length
                // buffer shares nothing, but the base chain carries owndata and the setflags(write=True)
                // base-writeability answer).
                var emptySlice = ArraySlice.Allocate(_typecode, 0, false);
                var emptyStorage = new UnmanagedStorage();
                emptyStorage._Allocate(slicedShape, emptySlice);
                emptyStorage._baseStorage = _baseStorage ?? this;
                emptyStorage.OnReshaped(); // a view never owns its buffer (NumPy: base != NULL ⟹ !OWNDATA)
                return emptyStorage;
            }

            // NumPy-aligned optimization: For contiguous slices, slice the InternalArray directly
            // and create a fresh shape with offset=0. This matches NumPy's data pointer adjustment
            // for contiguous views and makes IsSliced=false for contiguous slices.
            if (slicedShape.IsContiguous && slicedShape.offset > 0)
            {
                // Create a fresh contiguous shape (no offset)
                var freshShape = new Shape(slicedShape.dimensions);
                // A view of a non-writeable array (broadcast, or a read-only 'r' memmap) must stay
                // non-writeable — the fresh shape above defaults WRITEABLE back to true, which would let
                // a write reach read-only memory (a hard segfault on memory-mapped read-only pages, and
                // a silent NumPy-divergence for broadcasts). Shape.Slice already cleared it on
                // slicedShape; carry that through.
                if (!slicedShape.IsWriteable)
                    freshShape = freshShape.WithFlags(flagsToClear: ArrayFlags.WRITEABLE);
                var view = new UnmanagedStorage(InternalArray.Slice(slicedShape.offset, slicedShape.size), freshShape);
                view._baseStorage = _baseStorage ?? this;
                view.OnReshaped(); // a view never owns its buffer (NumPy: base != NULL ⟹ !OWNDATA)
                return view;
            }

            // Non-contiguous slices: create an alias with the sliced shape.
            // The slicedShape contains the correct offset and strides computed by Shape.Slice().
            // Views with non-zero offset or non-standard strides use coordinate-based access.
            return Alias(slicedShape);
        }

        private IEnumerable<Slice> ExpandEllipsis(Slice[] slices)
        {
            // count dimensions without counting ellipsis or newaxis
            var count = 0;
            foreach (var slice in slices)
            {
                if (slice.IsNewAxis || slice.IsEllipsis)
                    continue;
                count++;
            }

            // expand 
            foreach (var slice in slices)
            {
                if (slice.IsEllipsis)
                {
                    for (int i = 0; i < Shape.NDim - count; i++)
                        yield return Slice.All;
                    continue;
                }

                yield return slice;
            }
        }

        #endregion
    }
}