From b07d1dc6cac71b85bbc34cce0d7c5e90cf05ef16 Mon Sep 17 00:00:00 2001
From: Connor Ward <c.ward20@imperial.ac.uk>
Date: Thu, 19 Dec 2024 17:25:23 +0000
Subject: [PATCH] fixups

---
 firedrake/preconditioners/fdm.py   |  8 ++++++--
 firedrake/preconditioners/patch.py | 13 +++++++------
 firedrake/supermeshing.py          |  9 +++++----
 3 files changed, 18 insertions(+), 12 deletions(-)

diff --git a/firedrake/preconditioners/fdm.py b/firedrake/preconditioners/fdm.py
index 1696801cb4..186fbbc28c 100644
--- a/firedrake/preconditioners/fdm.py
+++ b/firedrake/preconditioners/fdm.py
@@ -1835,13 +1835,17 @@ def setSubMatCSR(comm, triu=False):
             return cache.setdefault(key, SparseAssembler.load_setSubMatCSR(comm, triu))
 
     @staticmethod
-    def load_c_code(code, name, **kwargs):
+    def load_c_code(code, name, comm, argtypes, restype):
         petsc_dir = get_petsc_dir()
         cppargs = [f"-I{d}/include" for d in petsc_dir]
         ldargs = ([f"-L{d}/lib" for d in petsc_dir]
                   + [f"-Wl,-rpath,{d}/lib" for d in petsc_dir]
                   + ["-lpetsc", "-lm"])
-        return load(code, "c", name, cppargs=cppargs, ldargs=ldargs, **kwargs)
+        dll = load(code, "c", cppargs=cppargs, ldargs=ldargs, comm=comm)
+        fn = getattr(dll, name)
+        fn.argtypes = argtypes
+        fn.restype = restype
+        return fn
 
     @staticmethod
     def load_setSubMatCSR(comm, triu=False):
diff --git a/firedrake/preconditioners/patch.py b/firedrake/preconditioners/patch.py
index 0a7bad5575..5e9d0d4fa0 100644
--- a/firedrake/preconditioners/patch.py
+++ b/firedrake/preconditioners/patch.py
@@ -505,12 +505,13 @@ def load_c_function(code, name, comm):
     ldargs = (["-L%s/lib" % d for d in get_petsc_dir()]
               + ["-Wl,-rpath,%s/lib" % d for d in get_petsc_dir()]
               + ["-lpetsc", "-lm"])
-    return load(code, "c", name,
-                argtypes=[ctypes.c_voidp, ctypes.c_int, ctypes.c_voidp,
-                          ctypes.c_voidp, ctypes.c_voidp, ctypes.c_int,
-                          ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp],
-                restype=ctypes.c_int, cppargs=cppargs, ldargs=ldargs,
-                comm=comm)
+    dll = load(code, "c", cppargs=cppargs, ldargs=ldargs, comm=comm)
+    fn = getattr(dll, name)
+    fn.argtypes = [ctypes.c_voidp, ctypes.c_int, ctypes.c_voidp,
+                   ctypes.c_voidp, ctypes.c_voidp, ctypes.c_int,
+                   ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp]
+    fn.restype = ctypes.c_int
+    return fn
 
 
 def make_c_arguments(form, kernel, state, get_map, require_state=False,
diff --git a/firedrake/supermeshing.py b/firedrake/supermeshing.py
index a1ce2cde17..6af10002d0 100644
--- a/firedrake/supermeshing.py
+++ b/firedrake/supermeshing.py
@@ -432,14 +432,15 @@ def likely(cell_A):
     includes = ["-I%s/include" % d for d in dirs]
     libs = ["-L%s/lib" % d for d in dirs]
     libs = libs + ["-Wl,-rpath,%s/lib" % d for d in dirs] + ["-lpetsc", "-lsupermesh"]
-    lib = load(
-        supermesh_kernel_str, "c", "supermesh_kernel",
+    dll = load(
+        supermesh_kernel_str, "c",
         cppargs=includes,
         ldargs=libs,
-        argtypes=[ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp],
-        restype=ctypes.c_int,
         comm=mesh_A._comm
     )
+    lib = getattr(dll, "supermesh_kernel")
+    lib.argtypes = [ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp]
+    lib.restype = ctypes.c_int
 
     ammm(V_A, V_B, likely, node_locations_A, node_locations_B, M_SS, ctypes.addressof(lib), mat)
     if orig_value_size == 1: