traits.h

// Copyright (C) 2024-2026 Joseph P. Dean and Garth N. Wells
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include <array>
#include <basix/mdspan.hpp>
#include <concepts>
#include <cstdint>
#include <dolfinx/common/types.h>
#include <ranges>
#include <span>
#include <type_traits>
 
namespace dolfinx::fem
{
template <class U, class T>
concept DofTransformKernel
    = std::is_invocable_v<U, std::span<T>, std::span<const std::uint32_t>,
                          std::int32_t, int>;

template <typename F>
constexpr bool is_transform_set(const F& fn)
{
  // A reference to a function is never null, and converting one to
  // bool warns under -Waddress rather than answering the question.
  if constexpr (std::is_function_v<F>)
    return true;
  else if constexpr (requires { static_cast<bool>(fn); })
    return static_cast<bool>(fn);
  else
    return true;
}

template <class U, class T, class G = dolfinx::scalar_value_t<T>>
concept FEkernel = std::is_invocable_v<U, T*, const T*, const T*, const G*,
                                       const int*, const std::uint8_t*, void*>;

template <class T>
concept MDSpan2
    = std::is_convertible_v<
          std::remove_cvref_t<T>,
          md::mdspan<const std::int32_t, md::dextents<std::size_t, 2>>>
      or std::is_convertible_v<
          std::remove_cvref_t<T>,
          md::mdspan<const std::int32_t, md::dextents<std::size_t, 1>>>;

template <class T>
concept MDSpan2Int32
    = dolfinx::MDSpanRank2<T>
      and std::same_as<typename std::remove_cvref_t<T>::value_type,
                       std::int32_t>;

template <class T, class U>
concept MDSpan2Floating
    = std::floating_point<U> and dolfinx::MDSpanRank2<T>
      and std::same_as<typename std::remove_cvref_t<T>::value_type, U>;

template <class C>
concept IndexList
    = std::ranges::random_access_range<C> and std::ranges::sized_range<C>
      and std::same_as<std::ranges::range_value_t<C>, std::int32_t>
      and requires(const std::remove_reference_t<C>& cells, std::size_t i) {
            { cells[i] } -> std::convertible_to<std::int32_t>;
          };

template <class D, std::floating_point U>
  requires MDSpan2Int32<D>
struct GeometryPack
{
  D dofmap;

  md::mdspan<const U, md::extents<std::size_t, md::dynamic_extent, 3>> x;
};

template <class D, class B, class E>
struct DofMapPack
{
  D map;

  B bs;

  E entities;
};

template <class T>
concept DofMapPackBase = requires(const std::remove_cvref_t<T>& t) {
  requires MDSpan2Int32<decltype(t.map)>;
  { t.bs } -> std::convertible_to<int>;
};

template <class T>
concept DofMapPackCells
    = DofMapPackBase<T> and requires(const std::remove_cvref_t<T>& t) {
        requires std::ranges::view<std::remove_cvref_t<decltype(t.entities)>>;
        requires IndexList<std::remove_cvref_t<decltype(t.entities)>>;
      };

template <class T>
concept DofMapPackEntities
    = DofMapPackBase<T> and requires(const std::remove_cvref_t<T>& t) {
        { t.entities(0, 0) } -> std::convertible_to<std::int32_t>;
      };

template <class T>
concept DofMapPackFacets
    = DofMapPackBase<T> and requires(const std::remove_cvref_t<T>& t) {
        { t.entities(0, 0, 0) } -> std::convertible_to<std::int32_t>;
      };

template <class P, class D, class B, class E>
struct FormArgument
{
  DofMapPack<D, B, E> dofmap;

  const P& transform;

  std::span<const std::uint32_t> cell_info;
};

template <class A, class T>
concept FormArgumentBase = requires(const std::remove_cvref_t<A>& a) {
  requires DofTransformKernel<std::remove_cvref_t<decltype(a.transform)>, T>;
  { a.cell_info } -> std::convertible_to<std::span<const std::uint32_t>>;
};

template <class A, class T>
concept FormArgumentCells
    = FormArgumentBase<A, T> and requires(const std::remove_cvref_t<A>& a) {
        requires DofMapPackCells<decltype(a.dofmap)>;
      };

template <class A, class T>
concept FormArgumentEntities
    = FormArgumentBase<A, T> and requires(const std::remove_cvref_t<A>& a) {
        requires DofMapPackEntities<decltype(a.dofmap)>;
      };

template <class A, class T>
concept FormArgumentFacets
    = FormArgumentBase<A, T> and requires(const std::remove_cvref_t<A>& a) {
        requires DofMapPackFacets<decltype(a.dofmap)>;
      };

template <class B>
inline constexpr bool is_scratch_buffer_type = false;
 
template <class T, std::size_t N>
inline constexpr bool is_scratch_buffer_type<std::array<T, N>> = true;
 
template <class T, std::size_t N>
inline constexpr bool is_scratch_buffer_type<std::span<T, N>> = true;

template <class B, class T>
concept ScratchBuffer
    = is_scratch_buffer_type<std::remove_cvref_t<B>>
      and std::ranges::contiguous_range<B> and std::ranges::output_range<B, T>
      and std::same_as<std::ranges::range_value_t<B>, T> and requires(B& b) {
            { b.data() } -> std::same_as<T*>;
            { b.size() } -> std::convertible_to<std::size_t>;
          };

template <class V, class T>
concept AssemblyVector
    = std::same_as<typename std::remove_cvref_t<V>::value_type, T>
      and requires(std::remove_cvref_t<V>& v, std::int32_t i) {
            { v[i] } -> std::convertible_to<T&>;
          };
 
namespace impl
{
using mdspan2_t = md::mdspan<const std::int32_t, md::dextents<std::size_t, 2>>;

template <class D>
inline constexpr std::size_t static_dofs_per_cell
    = std::remove_cvref_t<D>::static_extent(1);

template <class XD, std::floating_point U>
void gather_cell_coordinates(GeometryPack<XD, U> geometry, std::int32_t cell,
                             U* cdofs)
{
  const auto x_dofmap = geometry.dofmap;
  const std::size_t ndofs_x = x_dofmap.extent(1);
  const std::int32_t* xdofs
      = x_dofmap.data_handle() + static_cast<std::ptrdiff_t>(cell) * ndofs_x;
  const U* x = geometry.x.data_handle();
  for (std::size_t i = 0; i < ndofs_x; ++i)
  {
    const U* src = x + static_cast<std::ptrdiff_t>(xdofs[i]) * 3;
    for (std::size_t k = 0; k < 3; ++k)
      cdofs[3 * i + k] = src[k];
  }
}

template <class F>
void dispatch_bs(int bs, F&& f)
{
  switch (bs)
  {
  case 1:
    return f(std::integral_constant<int, 1>{});
  case 3:
    return f(std::integral_constant<int, 3>{});
  default:
    return f(bs);
  }
}

template <class F>
void dispatch_bs(int bs0, int bs1, F&& f)
{
  if (bs0 == bs1)
  {
    switch (bs0)
    {
    case 1:
      return f(std::integral_constant<int, 1>{},
               std::integral_constant<int, 1>{});
    case 3:
      return f(std::integral_constant<int, 3>{},
               std::integral_constant<int, 3>{});
    }
  }
 
  return f(bs0, bs1);
}
} // namespace impl
} // namespace dolfinx::fem