/* SPDX-License-Identifier: GPL-2.0-only */
/* Copyright (c) 2026 Meta Platforms, Inc. and affiliates. */

#ifndef T
#error "Define T (bit width: 32, 64) before including cnum_defs.h"
#endif

#include <linux/cnum.h>
#include <linux/kernel.h>
#include <linux/limits.h>
#include <linux/minmax.h>
#include <linux/compiler_types.h>

#define cnum_t   __PASTE(cnum, T)
#define ut       __PASTE(u, T)
#define st       __PASTE(s, T)
#define UT_MAX   __PASTE(__PASTE(U, T), _MAX)
#define ST_MAX   __PASTE(__PASTE(S, T), _MAX)
#define ST_MIN   __PASTE(__PASTE(S, T), _MIN)
#define EMPTY    __PASTE(__PASTE(CNUM, T), _EMPTY)
#define FN(name) __PASTE(__PASTE(cnum, T), __PASTE(_, name))

struct cnum_t FN(from_urange)(ut min, ut max)
{
        return (struct cnum_t){ .base = min, .size = (ut)max - min };
}

struct cnum_t FN(from_srange)(st min, st max)
{
        ut size = (ut)max - (ut)min;
        ut base = size == UT_MAX ? 0 : (ut)min;

        return (struct cnum_t){ .base = base, .size = size };
}

/* True if this cnum represents two unsigned ranges. */
static inline bool FN(urange_overflow)(struct cnum_t cnum)
{
        /* Same as cnum.base + cnum.size > UT_MAX but avoids overflow */
        return cnum.size > UT_MAX - (ut)cnum.base;
}

/*
 * cnum{T}_umin / cnum{T}_umax query an unsigned range represented by this cnum.
 * If cnum represents a range crossing the UT_MAX/0 boundary, the unbound range
 * [0..UT_MAX] is returned.
 */
ut FN(umin)(struct cnum_t cnum)
{
        return FN(urange_overflow)(cnum) ? 0 : cnum.base;
}
EXPORT_SYMBOL_GPL(FN(umin));

ut FN(umax)(struct cnum_t cnum)
{
        return FN(urange_overflow)(cnum) ? UT_MAX : cnum.base + cnum.size;
}
EXPORT_SYMBOL_GPL(FN(umax));

/* True if this cnum represents two signed ranges. */
static inline bool FN(srange_overflow)(struct cnum_t cnum)
{
        return FN(contains)(cnum, (ut)ST_MAX) && FN(contains)(cnum, (ut)ST_MIN);
}

/*
 * cnum{T}_smin / cnum{T}_smax query a signed range represented by this cnum.
 * If cnum represents a range crossing the ST_MAX/ST_MIN boundary, the unbound range
 * [ST_MIN..ST_MAX] is returned.
 */
st FN(smin)(struct cnum_t cnum)
{
        return FN(srange_overflow)(cnum)
               ? ST_MIN
               : min((st)cnum.base, (st)(cnum.base + cnum.size));
}

st FN(smax)(struct cnum_t cnum)
{
        return FN(srange_overflow)(cnum)
               ? ST_MAX
               : max((st)cnum.base, (st)(cnum.base + cnum.size));
}

/*
 * Returns a possibly empty intersection of cnums 'a' and 'b'.
 * If 'a' and 'b' intersect in two sub-arcs, the function over-approximates
 * and returns either 'a' or 'b', whichever is smaller.
 */
struct cnum_t FN(intersect)(struct cnum_t a, struct cnum_t b)
{
        struct cnum_t b1;
        ut dbase;

        if (FN(is_empty)(a) || FN(is_empty)(b))
                return EMPTY;

        if (a.base > b.base)
                swap(a, b);

        /*
         * Rotate frame of reference such that a.base is 0.
         * 'b1' is 'b' in this frame of reference.
         */
        dbase = b.base - a.base;
        b1 = (struct cnum_t){ dbase, b.size };
        if (FN(urange_overflow)(b1)) {
                if (b1.base <= a.size) {
                        /*
                         * Rotated frame (a.base at origin):
                         *
                         * 0                                       UT_MAX
                         * |--------------------------------------------|
                         * [=== a ==========================]           |
                         * [= b1 tail =]  [========= b1 main ==========>]
                         *                 ^-- b1.base <= a.size
                         *
                         * 'a' and 'b' intersect in two disjoint arcs,
                         * can't represent as single cnum, over-approximate
                         * the result.
                         */
                        return a.size <= b.size ? a : b;
                } else {
                        /*
                         * Rotated frame (a.base at origin):
                         *
                         * 0                                       UT_MAX
                         * |--------------------------------------------|
                         * [=== a =============]  |                     |
                         * [= b1 tail =]          [======= b1 main ====>]
                         *                         ^-- b1.base > a.size
                         *
                         * Only 'b' tail intersects 'a'.
                         */
                        return (struct cnum_t) {
                                .base = a.base,
                                .size = min(a.size, (ut)(b1.base + b1.size)),
                        };
                }
        } else if (a.size >= b1.base) {
                /*
                 * Rotated frame (a.base at origin):
                 *
                 * 0                                             UT_MAX
                 * |--------------------------------------------------|
                 * [=== a ==================================]         |
                 *                   [== b1 =====================]
                 *
                 * 0                                             UT_MAX
                 * |--------------------------------------------------|
                 * [=== a ==================================]         |
                 *                   [== b1 ====]
                 *                   ^-- b1.base <= a.size
                 *                   |<-- a.size - dbase -->|
                 *
                 * 'a' and 'b' intersect as one cnum.
                 */
                return (struct cnum_t) {
                        .base = b.base,
                        .size = min((ut)(a.size - dbase), b.size),
                };
        } else {
                return EMPTY;
        }
}

void FN(intersect_with)(struct cnum_t *dst, struct cnum_t src)
{
        *dst = FN(intersect)(*dst, src);
}

void FN(intersect_with_urange)(struct cnum_t *dst, ut min, ut max)
{
        FN(intersect_with)(dst, FN(from_urange)(min, max));
}

void FN(intersect_with_srange)(struct cnum_t *dst, st min, st max)
{
        FN(intersect_with)(dst, FN(from_srange)(min, max));
}

static inline struct cnum_t FN(normalize)(struct cnum_t cnum)
{
        if (cnum.size == UT_MAX && cnum.base != 0 && cnum.base != (ut)ST_MAX)
                cnum.base = 0;
        return cnum;
}

struct cnum_t FN(add)(struct cnum_t a, struct cnum_t b)
{
        if (FN(is_empty)(a) || FN(is_empty)(b))
                return EMPTY;
        if (a.size > UT_MAX - b.size)
                return (struct cnum_t){ 0, (ut)UT_MAX };
        else
                return FN(normalize)((struct cnum_t){ a.base + b.base, a.size + b.size });
}

struct cnum_t FN(negate)(struct cnum_t a)
{
        if (FN(is_empty)(a))
                return EMPTY;
        return FN(normalize)((struct cnum_t){ -((ut)a.base + a.size), a.size });
}

bool FN(is_empty)(struct cnum_t cnum)
{
        return cnum.base == EMPTY.base && cnum.size == EMPTY.size;
}

bool FN(contains)(struct cnum_t cnum, ut v)
{
        if (FN(is_empty)(cnum))
                return false;
        if (FN(urange_overflow)(cnum))
                return v >= cnum.base || v <= (ut)cnum.base + cnum.size;
        else
                return v >= cnum.base && v <= (ut)cnum.base + cnum.size;
}

bool FN(is_const)(struct cnum_t cnum)
{
        return cnum.size == 0;
}

bool FN(is_subset)(struct cnum_t bigger, struct cnum_t smaller)
{
        if (FN(is_empty(smaller)))
                return true;
        if (FN(is_empty(bigger)))
                return false;
        /* rotate both arcs such that 'bigger' starts at origin, hence does not overflow */
        smaller.base -= bigger.base;
        bigger.base = 0;
        if (FN(urange_overflow)(smaller) && bigger.size < UT_MAX)
                return false;
        return smaller.base + smaller.size <= bigger.size;
}

#undef EMPTY
#undef cnum_t
#undef ut
#undef st
#undef UT_MAX
#undef ST_MAX
#undef ST_MIN
#undef FN