nac3/nac3core/irrt/irrt_slice.hpp

#pragma once

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

namespace {
    // A proper slice in IRRT, all negative indices have be resolved to absolute values.
    // Even though nac3core's slices are always `int32_t`, we will template slice anyway
    // since this struct is used as a general utility.
    template <typename T>
    struct Slice {
        T start;
        T stop;
        T 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))`
        T len() {
            T 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;
            }
        }
    };

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

    // NOTE: using a bitfield for the `*_defined` is better, at the
    // cost of a more annoying implementation in nac3core inkwell
    template <typename T>
    struct UserSlice {
        uint8_t start_defined;
        T start;

        uint8_t stop_defined;
        T stop;

        uint8_t step_defined;
        T step;

        // Like Python's `slice(start, stop, step).indices(length)`
        Slice<T> indices(T 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<T> 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;
        }
    };
}