From 9c97dc47670323a843fd6ad2d40a476089d47045 Mon Sep 17 00:00:00 2001
From: "Garth N. Wells" <gnw20@cam.ac.uk>
Date: Thu, 7 Nov 2024 16:40:09 +0000
Subject: [PATCH] Small fixes

---
 cpp/dolfinx/fem/FunctionSpace.h | 62 ++++++++++++++++-----------------
 cpp/dolfinx/fem/utils.h         |  4 ++-
 2 files changed, 34 insertions(+), 32 deletions(-)

diff --git a/cpp/dolfinx/fem/FunctionSpace.h b/cpp/dolfinx/fem/FunctionSpace.h
index a0d77a283e..318670e96c 100644
--- a/cpp/dolfinx/fem/FunctionSpace.h
+++ b/cpp/dolfinx/fem/FunctionSpace.h
@@ -24,36 +24,6 @@
 
 namespace dolfinx::fem
 {
-/// @brief Compute the physical value shape of an element for a mesh.
-/// @param[in] element The element
-/// @param[in] tdim Topological dimension
-/// @param[in] gdim Geometric dimension
-/// @return Physical valus shape
-template <std::floating_point T>
-std::vector<std::size_t>
-compute_value_shape(const dolfinx::fem::FiniteElement<T>& element,
-                    std::size_t tdim, std::size_t gdim)
-{
-  std::vector<std::size_t> rvs = element.reference_value_shape();
-  std::vector<std::size_t> value_shape;
-  if (element.block_size() > 1)
-  {
-    value_shape = rvs;
-  }
-  else
-  {
-    for (std::size_t vs : rvs)
-    {
-      if (vs == tdim)
-        value_shape.push_back(gdim);
-      else
-        value_shape.push_back(vs);
-    }
-  }
-
-  return value_shape;
-}
-
 /// @brief This class represents a finite element function space defined
 /// by a mesh, a finite element, and a local-to-global map of the
 /// degrees-of-freedom.
@@ -70,7 +40,7 @@ class FunctionSpace
   /// @param[in] mesh Mesh that the space is defined on.
   /// @param[in] element Finite element for the space.
   /// @param[in] dofmap Degree-of-freedom map for the space.
-  /// @param[in] value_shape Shape of the value space on the physical cell.
+  /// @param[in] value_shape The shape of the value space on the physical cell
   FunctionSpace(std::shared_ptr<const mesh::Mesh<geometry_type>> mesh,
                 std::shared_ptr<const FiniteElement<geometry_type>> element,
                 std::shared_ptr<const DofMap> dofmap,
@@ -460,6 +430,36 @@ common_function_spaces(
   return {spaces0, spaces1};
 }
 
+/// @brief Compute the physical value shape of an element for a mesh
+/// @param[in] element The element
+/// @param[in] tdim Topological dimension
+/// @param[in] gdim Geometric dimension
+/// @return Physical valus shape
+template <std::floating_point T>
+std::vector<std::size_t>
+compute_value_shape(const dolfinx::fem::FiniteElement<T>& element,
+                    std::size_t tdim, std::size_t gdim)
+{
+  std::vector<std::size_t> rvs = element.reference_value_shape();
+  std::vector<std::size_t> value_shape(rvs.size());
+  if (element.block_size() > 1)
+  {
+    for (std::size_t i = 0; i < rvs.size(); ++i)
+      value_shape[i] = rvs[i];
+  }
+  else
+  {
+    for (std::size_t i = 0; i < rvs.size(); ++i)
+    {
+      if (rvs[i] == tdim)
+        value_shape[i] = gdim;
+      else
+        value_shape[i] = rvs[i];
+    }
+  }
+  return value_shape;
+}
+
 /// Type deduction
 template <typename U, typename V, typename W, typename X>
 FunctionSpace(U mesh, V element, W dofmap, X value_shape)
diff --git a/cpp/dolfinx/fem/utils.h b/cpp/dolfinx/fem/utils.h
index 880102306b..0a954d0e76 100644
--- a/cpp/dolfinx/fem/utils.h
+++ b/cpp/dolfinx/fem/utils.h
@@ -831,6 +831,7 @@ FunctionSpace<T> create_functionspace(
         reorder_fn
     = nullptr)
 {
+  assert(e);
   if (value_shape.has_value() and !e->reference_value_shape().empty())
   {
     throw std::runtime_error(
@@ -844,7 +845,7 @@ FunctionSpace<T> create_functionspace(
       = !value_shape.has_value() and !e->reference_value_shape().empty()
             ? fem::compute_value_shape(*e, mesh->topology()->dim(),
                                        mesh->geometry().dim())
-            : *value_shape;
+            : value_shape.value_or(std::vector<std::size_t>());
 
   // Create UFC subdofmaps and compute offset
   const int num_sub_elements = e->num_sub_elements();
@@ -871,6 +872,7 @@ FunctionSpace<T> create_functionspace(
   assert(mesh->topology());
   auto dofmap = std::make_shared<const DofMap>(create_dofmap(
       mesh->comm(), layout, *mesh->topology(), permute_inv, reorder_fn));
+
   return FunctionSpace(mesh, e, dofmap, _value_shape);
 }