Fix number of constants. (#665)

* Fix number of constants.
Resolves: FEniCS/dolfinx#3040

* Add test for constant renumbering (fails onmain)

* Apply suggestions from code review

Co-authored-by: Matthew Scroggs <matthew.w.scroggs@gmail.com>

---------

Co-authored-by: Matthew Scroggs <matthew.w.scroggs@gmail.com>

Author: jorgensd

ffcx/ir/representation.py

@@ -680,7 +680,7 @@ def _compute_expression_ir(expression, index, prefix, analysis, options, visuali
     # Build offsets for Constants
     original_constant_offsets = {}
     _offset = 0
-    for constant in ufl.algorithms.analysis.extract_constants(expression):
+    for constant in ufl.algorithms.analysis.extract_constants(original_expression):
         original_constant_offsets[constant] = _offset
         _offset += np.prod(constant.ufl_shape, dtype=int)
 

test/test_jit_expression.py

@@ -1,12 +1,13 @@
 # -*- coding: utf-8 -*-
-# Copyright (C) 2019-2022 Michal Habera and Jørgen S. Dokken
+# Copyright (C) 2019-2024 Michal Habera and Jørgen S. Dokken
 #
 # This file is part of FFCx.(https://www.fenicsproject.org)
 #
 # SPDX-License-Identifier:    LGPL-3.0-or-later
 
 import cffi
 import numpy as np
+import pytest
 
 import basix
 import basix.ufl
@@ -208,3 +209,45 @@ def exact_expr(x):
     exact = exact_expr(points.T)
 
     assert np.allclose(exact, output)
+
+
+def test_grad_constant(compile_args):
+    """Test if numbering of constants are correct after UFL eliminates the constant inside the gradient."""
+    c_el = basix.ufl.element("Lagrange", "triangle", 1, shape=(2, ))
+    mesh = ufl.Mesh(c_el)
+
+    x = ufl.SpatialCoordinate(mesh)
+    first_constant = ufl.Constant(mesh)
+    second_constant = ufl.Constant(mesh)
+    expr = second_constant * ufl.Dx(x[0]**2 + first_constant, 0)
+
+    dtype = np.float64
+    points = np.array([[0.33, 0.25]], dtype=dtype)
+
+    obj, _, _ = ffcx.codegeneration.jit.compile_expressions(
+        [(expr, points)], cffi_extra_compile_args=compile_args)
+
+    ffi = cffi.FFI()
+    expression = obj[0]
+
+    c_type = "double"
+    c_xtype = "double"
+
+    output = np.zeros(1, dtype=dtype)
+
+    # Define constants
+    coords = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=dtype)
+    u_coeffs = np.array([], dtype=dtype)
+    consts = np.array([3, 7], dtype=dtype)
+    entity_index = np.array([0], dtype=np.intc)
+    quad_perm = np.array([0], dtype=np.dtype("uint8"))
+
+    expression.tabulate_tensor_float64(
+        ffi.cast(f'{c_type} *', output.ctypes.data),
+        ffi.cast(f'{c_type} *', u_coeffs.ctypes.data),
+        ffi.cast(f'{c_type} *', consts.ctypes.data),
+        ffi.cast(f'{c_xtype} *', coords.ctypes.data),
+        ffi.cast('int *', entity_index.ctypes.data),
+        ffi.cast('uint8_t *', quad_perm.ctypes.data))
+
+    assert output[0] == pytest.approx(consts[1] * 2 * points[0, 0])