From fbab1ddd8765883253a7aec9c9419141f55242e3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20Schartum=20Dokken?= <dokken92@gmail.com>
Date: Tue, 27 Aug 2024 13:42:41 +0200
Subject: [PATCH] Fix DofMap::collapse() on rank with no dofs (#3369)

* Fix collapse on empty rank

* Add test that segfaults on main
---
 cpp/dolfinx/fem/DofMap.cpp          |  5 ++---
 python/test/unit/fem/test_dofmap.py | 23 +++++++++++++++++++++++
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/cpp/dolfinx/fem/DofMap.cpp b/cpp/dolfinx/fem/DofMap.cpp
index 871adce490d..5c10ac8d178 100644
--- a/cpp/dolfinx/fem/DofMap.cpp
+++ b/cpp/dolfinx/fem/DofMap.cpp
@@ -78,7 +78,8 @@ fem::DofMap build_collapsed_dofmap(const DofMap& dofmap_view,
   }
 
   // Create a map from old dofs to new dofs
-  std::vector<std::int32_t> old_to_new(dofs_view.back() + bs_view, -1);
+  std::size_t array_size = dofs_view.empty() ? 0 : dofs_view.back() + bs_view;
+  std::vector<std::int32_t> old_to_new(array_size, -1);
   for (std::size_t new_idx = 0; new_idx < sub_imap_to_imap.size(); ++new_idx)
   {
     for (int k = 0; k < bs_view; ++k)
@@ -215,7 +216,6 @@ std::pair<DofMap, std::vector<std::int32_t>> DofMap::collapse(
     reorder_fn = [](const graph::AdjacencyList<std::int32_t>& g)
     { return graph::reorder_gps(g); };
   }
-
   // Create new dofmap
   auto create_subdofmap = [](MPI_Comm comm, auto index_map_bs, auto& layout,
                              auto& topology, auto& reorder_fn, auto& dmap)
@@ -227,7 +227,6 @@ std::pair<DofMap, std::vector<std::int32_t>> DofMap::collapse(
 
       // Create new element dof layout and reset parent
       ElementDofLayout collapsed_dof_layout = layout.copy();
-
       auto [_index_map, bs, dofmaps] = build_dofmap_data(
           comm, topology, {collapsed_dof_layout}, reorder_fn);
       auto index_map
diff --git a/python/test/unit/fem/test_dofmap.py b/python/test/unit/fem/test_dofmap.py
index aa1ee0b4d21..85489d82fb0 100644
--- a/python/test/unit/fem/test_dofmap.py
+++ b/python/test/unit/fem/test_dofmap.py
@@ -421,3 +421,26 @@ def test_transpose_dofmap():
     dofmap = np.array([[0, 2, 1], [3, 2, 1], [4, 3, 1]], dtype=np.int32)
     transpose = dolfinx.fem.transpose_dofmap(dofmap, 3)
     assert np.array_equal(transpose.array, [0, 2, 5, 8, 1, 4, 3, 7, 6])
+
+
+def test_empty_rank_collapse():
+    """Test that dofmap with no dofs on a rank can be collapsed"""
+    if MPI.COMM_WORLD.rank == 0:
+        nodes = np.array([[0.0], [1.0], [2.0]], dtype=np.float64)
+        cells = np.array([[0, 1], [1, 2]], dtype=np.int64)
+    else:
+        nodes = np.empty((0, 1), dtype=np.float64)
+        cells = np.empty((0, 2), dtype=np.int64)
+    c_el = element("Lagrange", "interval", 1, shape=(1,))
+
+    def self_partitioner(comm: MPI.Intracomm, n, m, topo):
+        dests = np.full(len(topo[0]) // 2, comm.rank, dtype=np.int32)
+        offsets = np.arange(len(topo[0]) // 2 + 1, dtype=np.int32)
+        return dolfinx.graph.adjacencylist(dests, offsets)
+
+    mesh = create_mesh(MPI.COMM_WORLD, cells, nodes, c_el, partitioner=self_partitioner)
+
+    el = element("Lagrange", "interval", 1, shape=(2,))
+    V = functionspace(mesh, el)
+    V_0, _ = V.sub(0).collapse()
+    assert V.dofmap.index_map.size_local == V_0.dofmap.index_map.size_local