//==-------- alloc_util.hpp - Utilities for annotated usm allocation -------==//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#pragma once

#include <numeric>
#include <sycl/ext/intel/experimental/usm_properties.hpp>
#include <sycl/ext/oneapi/experimental/annotated_ptr/annotated_ptr.hpp>
#include <sycl/property_list.hpp>

namespace sycl {
inline namespace _V1 {
namespace ext {
namespace oneapi {
namespace experimental {

namespace detail {

////
//  Type traits for USM allocation with property support
////

// Check if a runtime property is valid
template <typename Prop> struct IsRuntimePropertyValid : std::false_type {};

// Validate the property list of annotated USM allocations by checking each
// property is
// 1. a valid compile-time property for annotated_ptr, or
// 2. a valid runtime property for annotated USM allocations
template <typename T, typename propertyList>
struct ValidAllocPropertyList : std::false_type {};
template <typename T>
struct ValidAllocPropertyList<T, empty_properties_t> : std::true_type {};
template <typename T, typename Prop, typename... Props>
struct ValidAllocPropertyList<T, detail::properties_t<Prop, Props...>>
    : std::integral_constant<
          bool, (is_valid_property<T *, Prop>::value ||
                 IsRuntimePropertyValid<Prop>::value) &&
                    ValidAllocPropertyList<
                        T, detail::properties_t<Props...>>::value> {
  static_assert(is_property_value_v<Prop>);
  static constexpr bool is_compile_time = std::is_empty_v<Prop>;
  // check if a compile-time property is valid for annotated_ptr
  static_assert(!is_compile_time || is_valid_property<T *, Prop>::value,
                "Found invalid compile-time property in the property list.");
  // check if a runtime property is valid for malloc
  static_assert(is_compile_time || IsRuntimePropertyValid<Prop>::value,
                "Found invalid runtime property in the property list.");
};

// The utility to filter a given property list to get the properties for
// annotated_ptr
template <typename PropertyListT> struct GetCompileTimeProperties {};

template <> struct GetCompileTimeProperties<empty_properties_t> {
  using type = empty_properties_t;
};

template <typename Prop>
struct GetCompileTimeProperties<detail::properties_t<Prop>> {
  using type =
      std::conditional_t<std::is_empty_v<Prop>, detail::properties_t<Prop>,
                         empty_properties_t>;
};

template <typename Prop, typename... Props>
struct GetCompileTimeProperties<detail::properties_t<Prop, Props...>> {
  using filtered_this_property_t =
      std::conditional_t<std::is_empty_v<Prop>, detail::properties_t<Prop>,
                         empty_properties_t>;
  using filtered_other_properties_t =
      typename GetCompileTimeProperties<detail::properties_t<Props...>>::type;
  using type = detail::merged_properties_t<filtered_this_property_t,
                                           filtered_other_properties_t>;
};

// Given the input property list `PropertyListT` and the usm kind `Kind` of
// annotated USM allocation functions, generate the property list for the
// returned annotated_ptr by following the rules:
// 1. all the compile-time properties in PropertyListT should appear in the
// output properties
// 2. if PropertyListT contains usm_kind, the value should be the same as
// `Kind`, usm_kind property should appear in the ouput properties
// 3. if PropertyListT does not contain usm_kind, a usm_kind property with value
// `Kind` is inserted in the ouput properties
template <sycl::usm::alloc Kind, typename PropertyListT>
struct GetAnnotatedPtrPropertiesWithUsmKind {};
template <sycl::usm::alloc Kind, typename... Props>
struct GetAnnotatedPtrPropertiesWithUsmKind<Kind,
                                            detail::properties_t<Props...>> {
  using input_properties_t = detail::properties_t<Props...>;
  using filtered_input_properties_t =
      typename GetCompileTimeProperties<input_properties_t>::type;

  static_assert(
      detail::get_property_or<usm_kind_key, input_properties_t>(usm_kind<Kind>)
              .value == Kind,
      "Input property list contains conflicting USM kind.");

  using type =
      detail::merged_properties_t<filtered_input_properties_t,
                                  decltype(properties{usm_kind<Kind>})>;
};

// Check if the 3 template parameters of annotated USM allocation functions,
// i.e.
//   T: allocated data type
//   propertyListA: input property list
//   propertyListB: property list for output anntoated_ptr
// meet the following conditions:
//  1. T is not a property list
//  2. propertyListB == all the compile-time properties in propertyListA
template <typename T, typename propertyListA, typename propertyListB>
struct CheckTAndPropLists : std::false_type {};

template <typename T, typename... PropsA, typename... PropsB>
struct CheckTAndPropLists<T, detail::properties_t<PropsA...>,
                          detail::properties_t<PropsB...>>
    : std::integral_constant<
          bool,
          !is_property_list<T>::value &&
              std::is_same_v<detail::properties_t<PropsB...>,
                             typename GetCompileTimeProperties<
                                 detail::properties_t<PropsA...>>::type>> {};

// Check if the 3 template parameters of annotated USM allocation functions of a
// certain usm kind:
//   T: allocated data type
//   propertyListA: input property list
//   propertyListB: property list for output anntoated_ptr
// meet the following conditions:
//  1. T is not a property list
//  2. if propertyListA contains usm_kind, the usm_kind must match with the
//  specified usm kind.
//     propertyListB is all the compile-time properties in propertyListA
//  3. if propertyListA does not contain usm_kind, propertyListB is all the
//  compile-time
//     properties in propertyListA with the usm_kind property inserted
template <sycl::usm::alloc Kind, typename T, typename propertyListA,
          typename propertyListB>
struct CheckTAndPropListsWithUsmKind : std::false_type {};

template <sycl::usm::alloc Kind, typename T, typename... PropsA,
          typename... PropsB>
struct CheckTAndPropListsWithUsmKind<Kind, T, detail::properties_t<PropsA...>,
                                     detail::properties_t<PropsB...>>
    : std::integral_constant<
          bool, !is_property_list<T>::value &&
                    std::is_same_v<
                        detail::properties_t<PropsB...>,
                        typename GetAnnotatedPtrPropertiesWithUsmKind<
                            Kind, detail::properties_t<PropsA...>>::type>> {};

} // namespace detail

////
//  Utility functions for USM allocation with property support
////

// Combine two alignment requirements for a pointer in the following way:
// 1. if either of the alignments is 0, return the other alignment
// 2. otherwise return the least common multiple of the two alignments
inline size_t combine_align(size_t alignA, size_t alignB) {
  return alignA == 0 ? alignB
                     : (alignB == 0 ? alignA : std::lcm(alignA, alignB));
}

} // namespace experimental
} // namespace oneapi
} // namespace ext
} // namespace _V1
} // namespace sycl
