Form.h

// Copyright (C) 2019-2026 Garth N. Wells, Chris Richardson, Joseph P. Dean and
// Jørgen S. Dokken
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "FunctionSpace.h"
#include "traits.h"
#include <algorithm>
#include <basix/mdspan.hpp>
#include <concepts>
#include <cstdint>
#include <dolfinx/common/types.h>
#include <dolfinx/mesh/EntityMap.h>
#include <dolfinx/mesh/Mesh.h>
#include <dolfinx/mesh/Topology.h>
#include <dolfinx/mesh/cell_types.h>
#include <format>
#include <functional>
#include <map>
#include <memory>
#include <ranges>
#include <set>
#include <span>
#include <stdexcept>
#include <tuple>
#include <utility>
#include <variant>
#include <vector>
 
namespace dolfinx::fem
{
template <dolfinx::scalar T>
class Constant;
template <dolfinx::scalar T, std::floating_point U>
class Function;

enum class IntegralType : std::int8_t
{
  cell = 0,           
  exterior_facet = 1, 
  interior_facet = 2, 
  vertex = 3,         
  ridge = 4           
};

constexpr int integral_entity_dim(IntegralType type, int tdim)
{
  switch (type)
  {
  case IntegralType::exterior_facet:
  case IntegralType::interior_facet:
    return tdim - 1;
  case IntegralType::ridge:
    return tdim - 2;
  case IntegralType::vertex:
    return 0;
  case IntegralType::cell:
    return tdim;
  }
 
  throw std::invalid_argument("Unknown integral type.");
}
 
namespace impl
{
inline md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>>
entity_permutations(mesh::Topology& topology, IntegralType type,
                    mesh::CellType cell_type)
{
  if (type == IntegralType::cell)
    return {};
 
  const int tdim = topology.dim();
  const int dim = integral_entity_dim(type, tdim);
 
  topology.create_entity_permutations(dim);
  const std::vector<std::uint8_t>& p = topology.get_entity_permutations(dim);
  const int num_entities_per_cell = mesh::cell_num_entities(cell_type, dim);
  return md::mdspan(p.data(), p.size() / num_entities_per_cell,
                    num_entities_per_cell);
}
} // namespace impl

template <dolfinx::scalar T, std::floating_point U = scalar_value_t<T>>
struct integral_data
{
  template <typename K, typename V, typename W>
    requires std::is_convertible_v<
                 std::remove_cvref_t<K>,
                 std::function<void(T*, const T*, const T*, const U*,
                                    const int*, const uint8_t*, void*)>>
                 and std::is_convertible_v<std::remove_cvref_t<V>,
                                           std::vector<std::int32_t>>
                 and std::is_convertible_v<std::remove_cvref_t<W>,
                                           std::vector<int>>
  integral_data(K&& kernel, V&& entities, W&& coeffs)
      : kernel(std::forward<K>(kernel)), entities(std::forward<V>(entities)),
        coeffs(std::forward<W>(coeffs))
  {
  }

  std::function<void(T*, const T*, const T*, const U*, const int*,
                     const uint8_t*, void*)>
      kernel;

  std::vector<std::int32_t> entities;

  std::vector<int> coeffs;
};

template <dolfinx::scalar T, std::floating_point U = dolfinx::scalar_value_t<T>>
class Form
{
public:
  using scalar_type = T;

  using geometry_type = U;

  template <typename X>
    requires std::is_convertible_v<
                 std::remove_cvref_t<X>,
                 std::map<std::tuple<IntegralType, int, int>,
                          integral_data<scalar_type, geometry_type>>>
  Form(
      const std::vector<std::shared_ptr<const FunctionSpace<geometry_type>>>& V,
      X&& integrals, std::shared_ptr<const mesh::Mesh<geometry_type>> mesh,
      const std::vector<
          std::shared_ptr<const Function<scalar_type, geometry_type>>>&
          coefficients,
      const std::vector<std::shared_ptr<const Constant<scalar_type>>>&
          constants,
      bool needs_facet_permutations,
      const std::vector<std::reference_wrapper<const mesh::EntityMap>>&
          entity_maps)
      : _function_spaces(V), _integrals(std::forward<X>(integrals)),
        _mesh(std::move(mesh)), _coefficients(coefficients),
        _constants(constants),
        _needs_facet_permutations(needs_facet_permutations)
  {
    if (!_mesh)
      throw std::invalid_argument("Form Mesh is null.");
 
    // `_mesh` is fixed for the remainder of construction, so its
    // topology and dimension are fetched once and reused below rather
    // than being re-fetched for every integral/coefficient.
    const mesh::Topology& topology = *_mesh->topology();
    const int tdim = topology.dim();
 
    // Map the integration entities of one integral to the
    // argument/coefficient domain, checking first that the mapping is
    // expressible: an integration entity maps to a *cell* of that mesh,
    // so unless the meshes have equal dimension the integral's entity
    // dimension must be that mesh's.
    auto map_entities
        = [tdim,
           &topology](IntegralType type, std::span<const std::int32_t> entities,
                      const mesh::Topology& topology0,
                      const mesh::EntityMap& emap) -> std::vector<std::int32_t>
    {
      if (type == IntegralType::cell)
      {
        return mesh::extract_cells_from_entities(
            topology0, topology, md::mdspan(entities.data(), entities.size()),
            std::cref(emap));
      }
 
      if (type == IntegralType::vertex)
      {
        throw std::invalid_argument(
            "Vertex integrals are not supported for a form with an argument "
            "or coefficient on another mesh. Supported types are cell, "
            "exterior facet, interior facet and ridge.");
      }
 
      const int dim0 = topology0.dim();
      const int edim = integral_entity_dim(type, tdim);
      if (tdim > dim0 and edim != dim0)
      {
        throw std::invalid_argument(std::format(
            "Cannot map integration entities of dimension {} to cells of a "
            "mesh of dimension {}. An argument or coefficient on another mesh "
            "must live on the entities being integrated over.",
            edim, dim0));
      }
 
      // Map the (cell, local_entity) pairs, flattened (interior facets
      // hold two pairs per entity), to cells of the argument/coefficient
      // mesh
      std::vector<std::int32_t> cells = mesh::extract_cells_from_entities(
          topology0, topology,
          md::mdspan<const std::int32_t,
                     md::extents<std::size_t, md::dynamic_extent, 2>>(
              entities.data(), entities.size() / 2, 2),
          std::cref(emap));
 
      // Replace the cell of each pair. Only the cell column is
      // meaningful: for codim > 0 the entity is itself the cell, so it
      // has no local index. The local index column is kept for the
      // layout used in packing/assembly.
      std::vector<std::int32_t> e(entities.begin(), entities.end());
      for (std::size_t i = 0; i < cells.size(); ++i)
        e[2 * i] = cells[i];
      return e;
    };
 
    _edata.reserve(_function_spaces.size());
    for (auto& space : _function_spaces)
    {
      // Working map: [integral type, integral_idx, kernel_idx]->entities
      std::map<std::tuple<IntegralType, int, int>,
               std::variant<std::vector<std::int32_t>,
                            std::span<const std::int32_t>>>
          vdata;
 
      if (auto mesh0 = space->mesh(); mesh0 == _mesh)
      {
        for (auto& [key, integral] : _integrals)
          vdata.insert({key, std::span(integral.entities)});
      }
      else
      {
        // Find correct entity map
        const mesh::Topology& topology0 = *mesh0->topology();
        const mesh::EntityMap& emap
            = mesh::find_entity_map(entity_maps, topology, topology0);
        for (auto& [key, itg] : _integrals)
        {
          auto [type, idx, kernel_idx] = key;
          vdata.insert(
              {key, map_entities(type, itg.entities, topology0, emap)});
        }
      }
 
      _edata.push_back(std::move(vdata));
    }
 
    for (auto& [key, integral] : _integrals)
    {
      auto [type, idx, kernel_idx] = key;
      for (int c : integral.coeffs)
      {
        if (auto mesh0 = coefficients.at(c)->function_space()->mesh();
            mesh0 == _mesh)
        {
          _cdata.insert({{type, idx, c}, std::span(integral.entities)});
        }
        else
        {
          // Find correct entity map
          const mesh::Topology& topology0 = *mesh0->topology();
          const mesh::EntityMap& emap
              = mesh::find_entity_map(entity_maps, topology, topology0);
          _cdata.insert(
              {{type, idx, c},
               map_entities(type, integral.entities, topology0, emap)});
        }
      }
    }
  }
 
  // Copy constructor (deleted). _edata and _cdata cache std::spans
  // aliasing the entity vectors owned by _integrals; a shallow copy
  // would leave the copy's spans pointing into the original's data.
  Form(const Form& form) = delete;

#ifdef _MSC_VER
  Form(Form&& form) noexcept = default;
#else
  Form(Form&& form) = default;
#endif

  ~Form() = default;
 
  // Copy assignment (deleted). Same aliasing reason as the copy
  // constructor.
  Form& operator=(const Form& form) = delete;

  Form& operator=(Form&& form) = default;

  int rank() const { return _function_spaces.size(); }

  std::shared_ptr<const mesh::Mesh<geometry_type>> mesh() const
  {
    return _mesh;
  }

  const std::vector<std::shared_ptr<const FunctionSpace<geometry_type>>>&
  function_spaces() const
  {
    return _function_spaces;
  }

  std::function<void(scalar_type*, const scalar_type*, const scalar_type*,
                     const geometry_type*, const int*, const uint8_t*, void*)>
  kernel(IntegralType type, int idx, int kernel_idx) const
  {
    auto it = _integrals.find({type, idx, kernel_idx});
    if (it == _integrals.end())
      throw std::out_of_range("Requested integral kernel not found.");
    return it->second.kernel;
  }

  std::set<IntegralType> integral_types() const
  {
    std::set<IntegralType> types;
    for (auto& [key, integral] : _integrals)
      types.insert(std::get<0>(key));
    return types;
  }

  std::vector<int> active_coeffs(IntegralType type, int idx) const
  {
    auto it = std::ranges::find_if(_integrals,
                                   [type, idx](auto& x)
                                   {
                                     auto [t, idx_, kernel_idx] = x.first;
                                     return t == type and idx_ == idx;
                                   });
    if (it == _integrals.end())
      throw std::out_of_range("Could not find active coefficient list.");
    return it->second.coeffs;
  }

  int num_integrals(IntegralType type, int kernel_idx) const
  {
    return std::ranges::count_if(_integrals,
                                 [type, kernel_idx](auto& x)
                                 {
                                   auto [t, id, k_idx] = x.first;
                                   return t == type and k_idx == kernel_idx;
                                 });
  }

  std::span<const std::int32_t> domain(IntegralType type, int idx,
                                       int kernel_idx) const
  {
    auto it = _integrals.find({type, idx, kernel_idx});
    if (it == _integrals.end())
      throw std::out_of_range("Requested domain not found.");
    return it->second.entities;
  }

  std::span<const std::int32_t> domain_arg(IntegralType type, int rank, int idx,
                                           int kernel_idx) const
  {
    auto it = _edata.at(rank).find({type, idx, kernel_idx});
    if (it == _edata.at(rank).end())
      throw std::out_of_range("Requested domain for argument not found.");
 
    return std::visit([](const auto& v) -> std::span<const std::int32_t>
                      { return v; }, it->second);
  }

  std::span<const std::int32_t> domain_coeff(IntegralType type, int idx,
                                             int c) const
  {
    auto it = _cdata.find({type, idx, c});
    if (it == _cdata.end())
      throw std::out_of_range("No domain for requested integral.");
    return std::visit([](const auto& v) -> std::span<const std::int32_t>
                      { return v; }, it->second);
  }

  const std::vector<
      std::shared_ptr<const Function<scalar_type, geometry_type>>>&
  coefficients() const
  {
    return _coefficients;
  }

  bool needs_facet_permutations() const { return _needs_facet_permutations; }

  std::vector<int> coefficient_offsets() const
  {
    std::vector<int> n{0};
    n.reserve(_coefficients.size() + 1);
    for (auto& c : _coefficients)
    {
      if (!c)
        throw std::runtime_error("Not all form coefficients have been set.");
      n.push_back(n.back() + c->function_space()->element()->space_dimension());
    }
    return n;
  }

  const std::vector<std::shared_ptr<const Constant<scalar_type>>>&
  constants() const
  {
    return _constants;
  }
 
private:
  // Function spaces (one for each argument)
  std::vector<std::shared_ptr<const FunctionSpace<geometry_type>>>
      _function_spaces;
 
  // Integrals (integral type, idx, kernel_idx)
  std::map<std::tuple<IntegralType, int, int>,
           integral_data<scalar_type, geometry_type>>
      _integrals;
 
  // The mesh
  std::shared_ptr<const mesh::Mesh<geometry_type>> _mesh;
 
  // Form coefficients
  std::vector<std::shared_ptr<const Function<scalar_type, geometry_type>>>
      _coefficients;
 
  // Constants associated with the Form
  std::vector<std::shared_ptr<const Constant<scalar_type>>> _constants;
 
  // True if permutation data needs to be passed into these integrals
  bool _needs_facet_permutations;
 
  // Mapped domain index data for argument functions.
  //
  // Consider:
  //
  // entities  = this->domain(type, idx, kernel_idx);
  // entities0 = _edata[0][{type, idx, kernel_idx}];
  //
  // Then `entities[i]` is a mesh entity index (e.g., cell index) in
  // `_mesh`, and  `entities0[i]` is the index of the same entity but in
  // the mesh associated with the argument 0 (test function) space.
  std::vector<std::map<
      std::tuple<IntegralType, int, int>,
      std::variant<std::vector<std::int32_t>, std::span<const std::int32_t>>>>
      _edata;
 
  // Mapped domain index data for coefficient functions.
  //
  // Consider:
  //
  // entities  = this->domain(type, idx, kernel_idx);
  // entities0 = _cdata[{type, idx, c}];
  //
  // where `c` is the coefficient index.
  //
  // Then `entities[i]` is a mesh entity index (e.g., cell index) in
  // `_mesh`, and  `entities0[i]` is the index of the same entity but in
  // the mesh associated with the coefficient Function.
  std::map<
      std::tuple<IntegralType, int, int>,
      std::variant<std::vector<std::int32_t>, std::span<const std::int32_t>>>
      _cdata;
};
} // namespace dolfinx::fem