`restrict=True` fails with RT/BDM/BDFM elements in parallel

Start with the reported reproducer and traceback, especially the assertion in firedrake/pyop3/insn/visitors.py:340 reached while assembling sparsity in firedrake/assemble.py. Compare the RT case with the working vector-valued CG case, then check whether BDM and BDFM reproduce it. Done means the reported parallel solve with restrict=True works for these elements.

Using restrict=True on an RT space in parallel fails. MFE:

from firedrake import *

mesh = UnitSquareMesh(4, 4)
V = FunctionSpace(mesh, "RT", 2)

u = TrialFunction(V)
v = TestFunction(V)
uh = Function(V)

a = inner(u, v) * dx
L = inner(Constant((1, 0)), v) * dx
bc = DirichletBC(V, Constant((0, 0)), "on_boundary")

solve(a == L, uh, bcs=bc, restrict=True)

I also tried BDM and BDFM elements and received the same error. A vector-valued CG space runs successfully however.

The traceback is

Traceback (most recent call last):
  File "/home/leo/work/firedrake-dev/fd-scripts/junk/restrict-rt-failure.py", line 14, in <module>
    solve(a == L, uh, bcs=bc, restrict=True)
    ~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "petsc4py/PETSc/Log.pyx", line 227, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "petsc4py/PETSc/Log.pyx", line 228, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "/home/leo/work/firedrake-dev/firedrake/firedrake/adjoint_utils/solving.py", line 57, in wrapper
    output = solve(*args, **kwargs)
  File "/home/leo/work/firedrake-dev/firedrake/firedrake/solving.py", line 160, in solve
    return _solve_varproblem(*args, **kwargs)
  File "/home/leo/work/firedrake-dev/firedrake/firedrake/solving.py", line 206, in _solve_varproblem
    solver = create_solver(problem, solver_parameters=solver_parameters,
                           nullspace=nullspace,
    ...<3 lines>...
                           appctx=appctx,
                           pre_apply_bcs=pre_apply_bcs)
  File "petsc4py/PETSc/Log.pyx", line 227, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "petsc4py/PETSc/Log.pyx", line 228, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "/nix/store/xdzsmlicnyznpa36hhk8fva9jp6bx49h-python3-3.13.14-env/lib/python3.13/contextlib.py", line 85, in inner
    return func(*args, **kwds)
  File "/home/leo/work/firedrake-dev/firedrake/firedrake/adjoint_utils/variational_solver.py", line 48, in wrapper
    init(self, problem, *args, **kwargs)
    ~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/leo/work/firedrake-dev/firedrake/firedrake/variational_solver.py", line 503, in __init__
    ctx.set_jacobian(self.snes)
    ~~~~~~~~~~~~~~~~^^^^^^^^^^^
  File "/home/leo/work/firedrake-dev/firedrake/firedrake/solving_utils.py", line 446, in set_jacobian
    snes.setJacobian(self.form_jacobian, J=self._jac.petscmat,
                                           ^^^^^^^^^
  File "/nix/store/xdzsmlicnyznpa36hhk8fva9jp6bx49h-python3-3.13.14-env/lib/python3.13/functools.py", line 1025, in __get__
    val = self.func(instance)
  File "/home/leo/work/firedrake-dev/firedrake/firedrake/solving_utils.py", line 723, in _jac
    return self._assembler_jac.allocate()
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^
  File "/home/leo/work/firedrake-dev/firedrake/pyop3/cache.py", line 801, in wrapper
    return func(*args, **kwargs)
  File "/home/leo/work/firedrake-dev/firedrake/firedrake/assemble.py", line 1303, in allocate
    sparsity = self._make_sparsity(
        test,
    ...<2 lines>...
        self._make_maps_and_regions(),
    )
  File "/home/leo/work/firedrake-dev/firedrake/firedrake/assemble.py", line 1379, in _make_sparsity
    op3.loop(
    ~~~~~~~~^
        loop_index,
        ^^^^^^^^^^^
        pack(sparsity[test_index_, trial_index_], loop_index, test_space, trial_space).assign(666, _weakref=False),
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        eager=True,
        ^^^^^^^^^^^
    )
    ^
  File "/home/leo/work/firedrake-dev/firedrake/pyop3/insn/base.py", line 997, in loop_
    return loop_expr(compiler_parameters=compiler_parameters) if eager else loop_expr
           ~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/leo/work/firedrake-dev/firedrake/pyop3/cache.py", line 801, in wrapper
    return func(*args, **kwargs)
  File "/home/leo/work/firedrake-dev/firedrake/pyop3/insn/base.py", line 126, in __call__
    self._get_execution_context(compiler_parameters)(**kwargs)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^
  File "/home/leo/work/firedrake-dev/firedrake/pyop3/insn/exec.py", line 162, in __call__
    executable = self.compile()
  File "/home/leo/work/firedrake-dev/.venv/real-xdzsmlicnyznpa36hhk8fva9jp6bx49h-python3-3.13.14-env/lib/python3.13/site-packages/cachetools/_cachedmethod.py", line 384, in __call__
    return wrapper(self._obj, *args, **kwargs)
  File "/home/leo/work/firedrake-dev/.venv/real-xdzsmlicnyznpa36hhk8fva9jp6bx49h-python3-3.13.14-env/lib/python3.13/site-packages/cachetools/_cachedmethod.py", line 367, in wrapper
    v = method(self, *args, **kwargs)
  File "/home/leo/work/firedrake-dev/firedrake/pyop3/insn/exec.py", line 183, in compile
    executor, orig_arg_index_to_buffer_map = self._compile()
                                             ~~~~~~~~~~~~~^^
  File "petsc4py/PETSc/Log.pyx", line 227, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "petsc4py/PETSc/Log.pyx", line 228, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "/home/leo/work/firedrake-dev/firedrake/pyop3/cache.py", line 641, in wrapper
    value = func(*args, **kwargs)
  File "/home/leo/work/firedrake-dev/firedrake/pyop3/insn/exec.py", line 244, in _compile
    self.preprocess()
    ~~~~~~~~~~~~~~~^^
  File "/home/leo/work/firedrake-dev/firedrake/pyop3/insn/exec.py", line 322, in preprocess
    insn = pyop3.insn.visitors.expand_transforms(insn)
  File "/nix/store/xdzsmlicnyznpa36hhk8fva9jp6bx49h-python3-3.13.14-env/lib/python3.13/functools.py", line 934, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^
  File "/home/leo/work/firedrake-dev/firedrake/pyop3/insn/visitors.py", line 164, in _
    return loop.record_new(statements=tuple(expand_transforms(s) for s in loop.statements))
                                      ~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/leo/work/firedrake-dev/firedrake/pyop3/insn/visitors.py", line 164, in <genexpr>
    return loop.record_new(statements=tuple(expand_transforms(s) for s in loop.statements))
                                            ~~~~~~~~~~~~~~~~~^^^
  File "/nix/store/xdzsmlicnyznpa36hhk8fva9jp6bx49h-python3-3.13.14-env/lib/python3.13/functools.py", line 934, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^
  File "/home/leo/work/firedrake-dev/firedrake/pyop3/insn/visitors.py", line 340, in _
    assert not any(
           ^^^^^^^^
    ...<2 lines>...
    )
    ^
AssertionError
--------------------------------------------------------------------------
MPI_ABORT was invoked on rank 1 in communicator MPI_COMM_WORLD
  Proc: [[41036,1],1]
  Errorcode: 1

NOTE: invoking MPI_ABORT causes Open MPI to kill all MPI processes.
You may or may not see output from other processes, depending on
exactly when Open MPI kills them.

This is running with Firedrake commit 2044cde1