nac3/nac3core/irrt/irrt_slice.hpp

#pragma once

#include "irrt_utils.hpp"
#include "irrt_typedefs.hpp"

namespace {
    struct Slice {
        SliceIndex start;
        SliceIndex stop;
        SliceIndex step;

        // The length/The number of elements of the slice if it were a range,
        // i.e., the value of `len(range(this->start, this->stop, this->end))`
        SliceIndex len() {
            SliceIndex diff = stop - start;
            if (diff > 0 && step > 0) {
                return ((diff - 1) / step) + 1;
            } else if (diff < 0 && step < 0) {
                return ((diff + 1) / step) + 1;
            } else {
                return 0;
            }
        }
    };

    SliceIndex resolve_index_in_length(SliceIndex length, SliceIndex index) {
        irrt_assert(length >= 0);
        if (index < 0) {
            // Remember that index is negative, so do a plus here
            return max<SliceIndex>(length + index, 0);
        } else {
            return min<SliceIndex>(length, index);
        }
    }

    // A user-written Python-like slice.
    //
    // i.e., this slice is a triple of either an int or nothing. (e.g., `my_array[:10:2]`, `start` is None)
    //
    // You can "resolve" a `UserSlice` by using `UserSlice::indices(<length>)`
    //
    // NOTE: using a bitfield for the `*_defined` is better, at the
    // cost of a more annoying implementation in nac3core inkwell
    struct UserSlice {
        // Did the user specify `start`? If 0, `start` is undefined (and contains an empty value)
        uint8_t start_defined;
        SliceIndex start;

        // Similar to `start_defined`
        uint8_t stop_defined;
        SliceIndex stop;

        // Similar to `start_defined`
        uint8_t step_defined;
        SliceIndex step;

        // Like Python's `slice(start, stop, step).indices(length)`
        Slice indices(SliceIndex length) {
            // NOTE: This function implements Python's `slice.indices` *FAITHFULLY*.
            // SEE: https://github.com/python/cpython/blob/f62161837e68c1c77961435f1b954412dd5c2b65/Objects/sliceobject.c#L546
            irrt_assert(length >= 0);
            irrt_assert(!step_defined || step != 0); // step_defined -> step != 0; step cannot be zero if specified by user

            Slice result;
            result.step = step_defined ? step : 1;
            bool step_is_negative = result.step < 0;

            if (start_defined) {
                result.start = resolve_index_in_length(length, start);
            } else {
                result.start = step_is_negative ? length - 1 : 0;
            }
            
            if (stop_defined) {
                result.stop = resolve_index_in_length(length, stop);
            } else {
                result.stop = step_is_negative ? -1 : length;
            }

            return result;
        }
    };
}