#pragma once

#include <c10/util/Exception.h>

namespace c10 {

namespace detail {
C10_API int64_t
maybe_wrap_dim_slow(int64_t dim, int64_t dim_post_expr, bool wrap_scalar);
}

static inline int64_t maybe_wrap_dim(
    int64_t dim,
    int64_t dim_post_expr,
    bool wrap_scalar = true) {
  // Inline the fast paths
  if (C10_LIKELY(-dim_post_expr <= dim && dim < dim_post_expr)) {
    // Branch-less version of dim + (dim < 0 ? dim_post_expr : 0)
    return dim + dim_post_expr * (dim < 0);
  }
  // Check edge-cases out-of-line (wrapping scalars and out-of-bounds errors)
  return c10::detail::maybe_wrap_dim_slow(dim, dim_post_expr, wrap_scalar);
}

} // namespace c10