Flatten intermediate representation of integral_names and subdomain_i…

…ds following reviewer comments.

ffcx/codegeneration/form.py

@@ -92,14 +92,13 @@ def generator(ir, options):
     for itg_type in ("cell", "exterior_facet", "interior_facet"):
         unsorted_integrals = []
         unsorted_ids = []
-        for key, name in ir.integral_names[itg_type].items():
-            for subdomain_id in ir.subdomain_ids[itg_type][key]:
-                unsorted_integrals += [L.AddressOf(L.Symbol(name))]
-                unsorted_ids += [subdomain_id]
+        for name, id in zip(ir.integral_names[itg_type], ir.subdomain_ids[itg_type], strict=True):
+            unsorted_integrals += [L.AddressOf(L.Symbol(name))]
+            unsorted_ids += [id]
 
         id_sort = numpy.argsort(unsorted_ids)
-        integrals += [unsorted_integrals[id_sort[i]] for i in range(len(unsorted_integrals))]
-        integral_ids += [unsorted_ids[id_sort[i]] for i in range(len(unsorted_integrals))]
+        integrals += [unsorted_integrals[i] for i in id_sort]
+        integral_ids += [unsorted_ids[i] for i in id_sort]
 
         integral_offsets.append(len(integrals))
 

ffcx/ir/representation.py

@@ -566,17 +566,19 @@ def _compute_form_ir(form_data, form_id, prefix, form_names, integral_names, ele
     ir["integral_names"] = {}
     ir["subdomain_ids"] = {}
     ufcx_integral_types = ("cell", "exterior_facet", "interior_facet")
-    ir["subdomain_ids"] = {itg_type: {} for itg_type in ufcx_integral_types}
-    ir["integral_names"] = {itg_type: {} for itg_type in ufcx_integral_types}
+    ir["subdomain_ids"] = {itg_type: [] for itg_type in ufcx_integral_types}
+    ir["integral_names"] = {itg_type: [] for itg_type in ufcx_integral_types}
     for itg_index, itg_data in enumerate(form_data.integral_data):
         # UFL is using "otherwise" for default integrals (over whole mesh)
         # but FFCx needs integers, so otherwise = -1
         integral_type = itg_data.integral_type
-        subdomain_ids = set(sid if sid != "otherwise" else -1 for sid in itg_data.subdomain_id)
+        subdomain_ids = [sid if sid != "otherwise" else -1 for sid in itg_data.subdomain_id]
+
         if min(subdomain_ids) < -1:
             raise ValueError("Integral subdomain IDs must be non-negative.")
-        ir["subdomain_ids"][integral_type][itg_index] = subdomain_ids
-        ir["integral_names"][integral_type][itg_index] = integral_names[(form_id, itg_index)]
+        ir["subdomain_ids"][integral_type] += subdomain_ids
+        for _ in range(len(subdomain_ids)):
+            ir["integral_names"][integral_type] += [integral_names[(form_id, itg_index)]]
 
     return FormIR(**ir)
 

test/test_jit_forms.py

@@ -899,7 +899,8 @@ def test_integral_grouping(compile_args):
     compiled_forms, module, _ = ffcx.codegeneration.jit.compile_forms(
         [a], cffi_extra_compile_args=compile_args)
     # NOTE: This assumes that the first integral type is cell integrals, see UFCx.h
-    num_integrals = compiled_forms[0].form_integral_offsets[1] - compiled_forms[0].form_integral_offsets[0]
+    cell = module.lib.cell
+    num_integrals = compiled_forms[0].form_integral_offsets[cell+1] - compiled_forms[0].form_integral_offsets[cell]
     assert num_integrals == 4
-    unique_integrals = set([compiled_forms[0].form_integrals[i] for i in range(num_integrals)])
+    unique_integrals = set([compiled_forms[0].form_integrals[compiled_forms[0].form_integral_offsets[cell]+i] for i in range(num_integrals)])
     assert len(unique_integrals) == 2