refactor: add builder.stash_integrals()

Author: ksagiyam

tsfc/driver.py

@@ -1,4 +1,3 @@
-import collections
 import operator
 import string
 import time
@@ -157,15 +156,11 @@ def compile_integral(integral_data, form_data, prefix, parameters, interface, co
         params = parameters.copy()
         params.update(integral.metadata())  # integral metadata overrides
 
-        mode = pick_mode(params["mode"])
-        mode_irs.setdefault(mode, collections.OrderedDict())
-
         integrand = ufl.replace(integral.integrand(), form_data.function_replace_map)
         integrand_exprs = builder.compile_ufl(integrand, params, ctx)
         integral_exprs = builder.construct_integrals(integrand_exprs, params)
+        builder.stash_integrals(integral_exprs, params, ctx)
         quad_rule = params["quadrature_rule"]
-        for var, rep in zip(return_variables, integral_exprs):
-            mode_irs[mode].setdefault(var, []).append(rep)
 
     # Finalise mode representations into a set of assignments
     assignments = []

tsfc/kernel_interface/common.py

@@ -149,6 +149,13 @@ def construct_integrals(self, expressions, params):
                               self.argument_multiindices,
                               params)
 
+    def stash_integrals(self, reps, params, ctx):
+        mode = pick_mode(params["mode"])
+        mode_irs = ctx['mode_irs']
+        mode_irs.setdefault(mode, collections.OrderedDict())
+        for var, rep in zip(self.return_variables, reps):
+            mode_irs[mode].setdefault(var, []).append(rep)
+
     @cached_property
     def fem_config(self):
         info = self.integral_data_info