Refactor JIT cache layer with JITCallbackSet

[lmoresi]

Summary

Builds on PR #92 (@gthyagi) which fixed the Integral/BdIntegral JIT cache collision bug. This PR refactors the fix into a reusable architectural pattern that simplifies the JIT interface for all callers.

What PR #92 identified

When Integral(fn=1.0) and BdIntegral(fn=1.0) ran on the same mesh, the JIT cache key was a flat tuple of all expressions — losing which callback role each expression belonged to. The second call reused the first's cached module, wiring the wrong PETSc callback slot and returning 0.0.

What this PR does

Introduces JITCallbackSet, a frozen dataclass that groups the five PETSc callback lists (residual, bcs, jacobian, bd_residual, bd_jacobian) into a single structured container.

Before (PR #92 fix): The cache key was made structured by manually building parallel tuples for each role inside getext(). Callers still passed 5 positional lists:

# Integral — which positional arg is which?
getext(mesh, [self.fn,], [], [], [], [], mesh.vars.values())
# BdIntegral
getext(mesh, [], [], [], [self.fn,], [], mesh.vars.values())

After: Callers declare intent by name, with a clean 3-argument interface:

# Integral — self-documenting
getext(mesh, JITCallbackSet(residual=(self.fn,)), mesh.vars.values())
# BdIntegral
getext(mesh, JITCallbackSet(bd_residual=(self.fn,)), mesh.vars.values())

The structured cache key is now a data structure invariant — you can't construct a callback set without declaring roles, so the collision bug cannot recur.

Changes

• JITCallbackSet dataclass with flat(), signature(), map(), counts
• prepare_for_cache_key() — extracted as a module-level function (was an inline closure). Two-phase expression preparation: constants → placeholders, then unwrap to pure SymPy. Reusable by any code needing structural expansion for hashing.
• getext(mesh, callbacks, primary_field_list, ...) — clean 3-positional-arg signature
• _createext(name, mesh, callbacks, primary_field_list, ...) — takes JITCallbackSet directly, uses callbacks.counts instead of len(fns_residual) etc.
• All 7 call sites updated: 3 solvers, 3 integrals, 1 test file

Files changed

File	Change
src/underworld3/utilities/_jitextension.py	Add JITCallbackSet, prepare_for_cache_key(); refactor getext() and _createext()
src/underworld3/cython/petsc_maths.pyx	Integral/BdIntegral use JITCallbackSet
src/underworld3/cython/petsc_generic_snes_solvers.pyx	Scalar/Vector/Stokes solvers use JITCallbackSet
tests/test_0004_pointwise_fns.py	Update direct getext() calls
tests/test_0502_boundary_integrals.py	PR #92 regression tests (spherical shell cache collision)

Test plan

• 21/21 boundary integral tests pass (including PR Fix Integral/BdIntegral JIT cache collisions #92 spherical shell regression)
• 18/18 core solver tests pass (Poisson, Stokes, VE-Stokes)
• 374/374 level_1 tests pass (0 failures, 7 skipped, 1 xfail)

Underworld development team with AI support from Claude Code

[lmoresi]

This is just some tidying prompted by @gthyagi 's bug fix. Hardening the JIT infrastructure is always and improvement.

[Copilot]

Pull request overview

Refactors the JIT extension cache interface to prevent callback-role collisions by introducing a structured callback container (JITCallbackSet) and updating all callers (integrals, solvers, and tests) to use it.

Changes:

• Introduces JITCallbackSet and refactors getext() / _createext() to accept structured callbacks (fixing role-collision cache bugs by construction).
• Extracts cache-key preparation into prepare_for_cache_key() and uses role-preserving signatures for hashing.
• Updates solver/integral call sites and adjusts tests, including spherical-shell regression coverage for boundary/volume cache collisions.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
src/underworld3/utilities/_jitextension.py	Adds JITCallbackSet, factors cache-key preparation, refactors getext()/_createext() to use structured callbacks and role-preserving signatures.
src/underworld3/cython/petsc_maths.pyx	Updates Integral / CellWiseIntegral / BdIntegral to call getext() with JITCallbackSet.
src/underworld3/cython/petsc_generic_snes_solvers.pyx	Updates SNES solvers to pass structured callbacks into getext().
tests/test_0004_pointwise_fns.py	Updates direct getext() test calls to use JITCallbackSet.
tests/test_0502_boundary_integrals.py	Adds/updates regression tests covering spherical-shell boundary/volume integral cache-collision scenarios.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

JITCallbackSet claims slots are tuples and getext() relies on tuple concatenation + hashing (signature()). If a caller passes a list (likely, since the old getext API accepted lists), flat()/signature() will raise (list+tuple TypeError or unhashable list). Consider coercing each slot to tuple (and treating None as empty) in post_init using object.setattr so the container is always immutable + hashable.

[Copilot]

The typing imports include List and Tuple, but this module no longer uses them (Optional is still used). Consider removing unused imports to avoid confusion and keep linters clean.

Suggested change

	from typing import List, Optional, Tuple
	from typing import Optional

[lmoresi]

Copilot fixes adopted. I think we can merge this as this passes all CI tests and it is really quite a simple task.