From 53c12e2dd7286e29fdc7ec06b197bdaad1ffb779 Mon Sep 17 00:00:00 2001
From: Connor Ward <c.ward20@imperial.ac.uk>
Date: Thu, 29 Jun 2023 14:26:10 +0100
Subject: [PATCH 1/3] Add Counted mixin class and refactor form signature
 computation

---
 ufl/algorithms/signature.py |  8 +----
 ufl/coefficient.py          | 12 ++------
 ufl/constant.py             | 10 ++-----
 ufl/core/multiindex.py      | 11 ++-----
 ufl/form.py                 | 58 ++++++++++++++++++++++++++++++-------
 ufl/matrix.py               | 10 ++-----
 ufl/utils/counted.py        | 40 +++++++++----------------
 ufl/variable.py             | 14 ++++-----
 8 files changed, 80 insertions(+), 83 deletions(-)

diff --git a/ufl/algorithms/signature.py b/ufl/algorithms/signature.py
index ab9690bd6..4650818c1 100644
--- a/ufl/algorithms/signature.py
+++ b/ufl/algorithms/signature.py
@@ -43,7 +43,6 @@ def compute_terminal_hashdata(expressions, renumbering):
     # arguments, and just take repr of the rest of the terminals while
     # we're iterating over them
     terminal_hashdata = {}
-    labels = {}
     index_numbering = {}
     for expression in expressions:
         for expr in traverse_unique_terminals(expression):
@@ -69,12 +68,7 @@ def compute_terminal_hashdata(expressions, renumbering):
                 data = expr._ufl_signature_data_(renumbering)
 
             elif isinstance(expr, Label):
-                # Numbering labels as we visit them # TODO: Include in
-                # renumbering
-                data = labels.get(expr)
-                if data is None:
-                    data = "L%d" % len(labels)
-                    labels[expr] = data
+                data = expr._ufl_signature_data_(renumbering)
 
             elif isinstance(expr, ExprList):
                 # Not really a terminal but can have 0 operands...
diff --git a/ufl/coefficient.py b/ufl/coefficient.py
index 51a508c18..61d2d2bad 100644
--- a/ufl/coefficient.py
+++ b/ufl/coefficient.py
@@ -19,13 +19,13 @@
 from ufl.functionspace import AbstractFunctionSpace, FunctionSpace, MixedFunctionSpace
 from ufl.form import BaseForm
 from ufl.split_functions import split
-from ufl.utils.counted import counted_init
+from ufl.utils.counted import Counted
 from ufl.duals import is_primal, is_dual
 
 # --- The Coefficient class represents a coefficient in a form ---
 
 
-class BaseCoefficient(object):
+class BaseCoefficient(Counted):
     """UFL form argument type: Parent Representation of a form coefficient."""
 
     # Slots are disabled here because they cause trouble in PyDOLFIN
@@ -33,14 +33,13 @@ class BaseCoefficient(object):
     # __slots__ = ("_count", "_ufl_function_space", "_repr", "_ufl_shape")
     _ufl_noslots_ = True
     __slots__ = ()
-    _globalcount = 0
     _ufl_is_abstract_ = True
 
     def __getnewargs__(self):
         return (self._ufl_function_space, self._count)
 
     def __init__(self, function_space, count=None):
-        counted_init(self, count, Coefficient)
+        Counted.__init__(self, count)
 
         if isinstance(function_space, FiniteElementBase):
             # For legacy support for .ufl files using cells, we map
@@ -57,9 +56,6 @@ def __init__(self, function_space, count=None):
         self._repr = "BaseCoefficient(%s, %s)" % (
             repr(self._ufl_function_space), repr(self._count))
 
-    def count(self):
-        return self._count
-
     @property
     def ufl_shape(self):
         "Return the associated UFL shape."
@@ -119,7 +115,6 @@ class Cofunction(BaseCoefficient, BaseForm):
         "_ufl_shape",
         "_hash"
     )
-    # _globalcount = 0
     _primal = False
     _dual = True
 
@@ -162,7 +157,6 @@ class Coefficient(FormArgument, BaseCoefficient):
     """UFL form argument type: Representation of a form coefficient."""
 
     _ufl_noslots_ = True
-    _globalcount = 0
     _primal = True
     _dual = False
 
diff --git a/ufl/constant.py b/ufl/constant.py
index f2359cc7c..704604f7a 100644
--- a/ufl/constant.py
+++ b/ufl/constant.py
@@ -11,17 +11,16 @@
 from ufl.core.ufl_type import ufl_type
 from ufl.core.terminal import Terminal
 from ufl.domain import as_domain
-from ufl.utils.counted import counted_init
+from ufl.utils.counted import Counted
 
 
 @ufl_type()
-class Constant(Terminal):
+class Constant(Terminal, Counted):
     _ufl_noslots_ = True
-    _globalcount = 0
 
     def __init__(self, domain, shape=(), count=None):
         Terminal.__init__(self)
-        counted_init(self, count=count, countedclass=Constant)
+        Counted.__init__(self, count)
 
         self._ufl_domain = as_domain(domain)
         self._ufl_shape = shape
@@ -31,9 +30,6 @@ def __init__(self, domain, shape=(), count=None):
         self._repr = "Constant({}, {}, {})".format(
             repr(self._ufl_domain), repr(self._ufl_shape), repr(self._count))
 
-    def count(self):
-        return self._count
-
     @property
     def ufl_shape(self):
         return self._ufl_shape
diff --git a/ufl/core/multiindex.py b/ufl/core/multiindex.py
index 55857b8c7..5ae353703 100644
--- a/ufl/core/multiindex.py
+++ b/ufl/core/multiindex.py
@@ -10,7 +10,7 @@
 # Modified by Massimiliano Leoni, 2016.
 
 
-from ufl.utils.counted import counted_init
+from ufl.utils.counted import Counted
 from ufl.core.ufl_type import ufl_type
 from ufl.core.terminal import Terminal
 
@@ -70,20 +70,15 @@ def __repr__(self):
         return r
 
 
-class Index(IndexBase):
+class Index(IndexBase, Counted):
     """UFL value: An index with no value assigned.
 
     Used to represent free indices in Einstein indexing notation."""
     __slots__ = ("_count",)
 
-    _globalcount = 0
-
     def __init__(self, count=None):
         IndexBase.__init__(self)
-        counted_init(self, count, Index)
-
-    def count(self):
-        return self._count
+        Counted.__init__(self, count)
 
     def __hash__(self):
         return hash(("Index", self._count))
diff --git a/ufl/form.py b/ufl/form.py
index 9c4ab0eda..5124d9d13 100644
--- a/ufl/form.py
+++ b/ufl/form.py
@@ -16,12 +16,15 @@
 from itertools import chain
 
 from ufl.checks import is_scalar_constant_expression
+from ufl.constant import Constant
 from ufl.constantvalue import Zero
 from ufl.core.expr import Expr, ufl_err_str
 from ufl.core.ufl_type import UFLType, ufl_type
 from ufl.domain import extract_unique_domain, sort_domains
 from ufl.equation import Equation
 from ufl.integral import Integral
+from ufl.utils.counted import Counted
+from ufl.utils.sorting import sorted_by_count
 
 # Export list for ufl.classes
 __all_classes__ = ["Form", "BaseForm", "ZeroBaseForm"]
@@ -257,8 +260,9 @@ class Form(BaseForm):
         "_arguments",
         "_coefficients",
         "_coefficient_numbering",
-        "_constant_numbering",
         "_constants",
+        "_constant_numbering",
+        "_terminal_numbering",
         "_hash",
         "_signature",
         # --- Dict that external frameworks can place framework-specific
@@ -289,11 +293,10 @@ def __init__(self, integrals):
         self._coefficients = None
         self._coefficient_numbering = None
         self._constant_numbering = None
+        self._terminal_numbering = None
 
         from ufl.algorithms.analysis import extract_constants
         self._constants = extract_constants(self)
-        self._constant_numbering = dict(
-            (c, i) for i, c in enumerate(self._constants))
 
         # Internal variables for caching of hash and signature after
         # first request
@@ -406,8 +409,15 @@ def coefficients(self):
     def coefficient_numbering(self):
         """Return a contiguous numbering of coefficients in a mapping
         ``{coefficient:number}``."""
+        # cyclic import
+        from ufl.coefficient import Coefficient
+
         if self._coefficient_numbering is None:
-            self._analyze_form_arguments()
+            self._coefficient_numbering = {
+                expr: num
+                for expr, num in self.terminal_numbering().items()
+                if isinstance(expr, Coefficient)
+            }
         return self._coefficient_numbering
 
     def constants(self):
@@ -416,8 +426,38 @@ def constants(self):
     def constant_numbering(self):
         """Return a contiguous numbering of constants in a mapping
         ``{constant:number}``."""
+        if self._constant_numbering is None:
+            self._constant_numbering = {
+                expr: num
+                for expr, num in self.terminal_numbering().items()
+                if isinstance(expr, Constant)
+            }
         return self._constant_numbering
 
+    def terminal_numbering(self):
+        """Return a contiguous numbering for all counted objects in the form.
+
+        The returned object is mapping from terminal to its number (an integer).
+
+        The numbering is computed per type so :class:`Coefficient`s,
+        :class:`Constant`s, etc will each be numbered from zero.
+
+        """
+        # cyclic import
+        from ufl.algorithms.analysis import extract_type
+
+        if self._terminal_numbering is None:
+            exprs_by_type = defaultdict(set)
+            for counted_expr in extract_type(self, Counted):
+                exprs_by_type[type(counted_expr)].add(counted_expr)
+
+            numbering = {}
+            for type_, exprs in exprs_by_type.items():
+                for i, expr in enumerate(sorted_by_count(exprs)):
+                    numbering[expr] = i
+            self._terminal_numbering = numbering
+        return self._terminal_numbering
+
     def signature(self):
         "Signature for use with jit cache (independent of incidental numbering of indices etc.)"
         if self._signature is None:
@@ -625,23 +665,19 @@ def _analyze_form_arguments(self):
             sorted(set(arguments), key=lambda x: x.number()))
         self._coefficients = tuple(
             sorted(set(coefficients), key=lambda x: x.count()))
-        self._coefficient_numbering = dict(
-            (c, i) for i, c in enumerate(self._coefficients))
 
     def _compute_renumbering(self):
         # Include integration domains and coefficients in renumbering
         dn = self.domain_numbering()
-        cn = self.coefficient_numbering()
-        cnstn = self.constant_numbering()
+        tn = self.terminal_numbering()
         renumbering = {}
         renumbering.update(dn)
-        renumbering.update(cn)
-        renumbering.update(cnstn)
+        renumbering.update(tn)
 
         # Add domains of coefficients, these may include domains not
         # among integration domains
         k = len(dn)
-        for c in cn:
+        for c in self.coefficients():
             d = extract_unique_domain(c)
             if d is not None and d not in renumbering:
                 renumbering[d] = k
diff --git a/ufl/matrix.py b/ufl/matrix.py
index 58e704d2b..498b2fb5e 100644
--- a/ufl/matrix.py
+++ b/ufl/matrix.py
@@ -13,13 +13,13 @@
 from ufl.core.ufl_type import ufl_type
 from ufl.argument import Argument
 from ufl.functionspace import AbstractFunctionSpace
-from ufl.utils.counted import counted_init
+from ufl.utils.counted import Counted
 
 
 # --- The Matrix class represents a matrix, an assembled two form ---
 
 @ufl_type()
-class Matrix(BaseForm):
+class Matrix(BaseForm, Counted):
     """An assemble linear operator between two function spaces."""
 
     __slots__ = (
@@ -30,7 +30,6 @@ class Matrix(BaseForm):
         "_hash",
         "_ufl_shape",
         "_arguments")
-    _globalcount = 0
 
     def __getnewargs__(self):
         return (self._ufl_function_spaces[0], self._ufl_function_spaces[1],
@@ -38,7 +37,7 @@ def __getnewargs__(self):
 
     def __init__(self, row_space, column_space, count=None):
         BaseForm.__init__(self)
-        counted_init(self, count, Matrix)
+        Counted.__init__(self, count)
 
         if not isinstance(row_space, AbstractFunctionSpace):
             raise ValueError("Expecting a FunctionSpace as the row space.")
@@ -55,9 +54,6 @@ def __init__(self, row_space, column_space, count=None):
             repr(self._ufl_function_spaces[1]), repr(self._count)
         )
 
-    def count(self):
-        return self._count
-
     def ufl_function_spaces(self):
         "Get the tuple of function spaces of this coefficient."
         return self._ufl_function_spaces
diff --git a/ufl/utils/counted.py b/ufl/utils/counted.py
index 66d3fd79f..8528b30d6 100644
--- a/ufl/utils/counted.py
+++ b/ufl/utils/counted.py
@@ -1,5 +1,5 @@
 # -*- coding: utf-8 -*-
-"Utilites for types with a global unique counter attached to each object."
+"Mixin class for types with a global unique counter attached to each object."
 
 # Copyright (C) 2008-2016 Martin Sandve Alnæs
 #
@@ -7,37 +7,25 @@
 #
 # SPDX-License-Identifier:    LGPL-3.0-or-later
 
+import itertools
 
-def counted_init(self, count=None, countedclass=None):
-    "Initialize a counted object, see ExampleCounted below for how to use."
 
-    if countedclass is None:
-        countedclass = type(self)
+class Counted:
+    """Mixin class for globally counted objects."""
 
-    if count is None:
-        count = countedclass._globalcount
+    # Mixin classes do not work well with __slots__ so _count must be
+    # added to the __slots__ of the inheriting class
+    __slots__ = ()
 
-    self._count = count
+    _counter = None
 
-    if self._count >= countedclass._globalcount:
-        countedclass._globalcount = self._count + 1
-
-
-class ExampleCounted(object):
-    """An example class for classes of objects identified by a global counter.
-
-    Mimic this class to create globally counted objects within a single type.
-    """
-    # Store the count for each object
-    __slots__ = ("_count",)
-
-    # Store a global counter with the class
-    _globalcount = 0
-
-    # Call counted_init with an optional constructor argument and the class
     def __init__(self, count=None):
-        counted_init(self, count, ExampleCounted)
+        # create a new counter for each subclass
+        cls = type(self)
+        if cls._counter is None:
+            cls._counter = itertools.count()
+
+        self._count = count if count is not None else next(cls._counter)
 
-    # Make the count accessible
     def count(self):
         return self._count
diff --git a/ufl/variable.py b/ufl/variable.py
index fb14c43d6..10e755b64 100644
--- a/ufl/variable.py
+++ b/ufl/variable.py
@@ -8,7 +8,7 @@
 #
 # SPDX-License-Identifier:    LGPL-3.0-or-later
 
-from ufl.utils.counted import counted_init
+from ufl.utils.counted import Counted
 from ufl.core.expr import Expr
 from ufl.core.ufl_type import ufl_type
 from ufl.core.terminal import Terminal
@@ -17,17 +17,12 @@
 
 
 @ufl_type()
-class Label(Terminal):
+class Label(Terminal, Counted):
     __slots__ = ("_count",)
 
-    _globalcount = 0
-
     def __init__(self, count=None):
         Terminal.__init__(self)
-        counted_init(self, count, Label)
-
-    def count(self):
-        return self._count
+        Counted.__init__(self, count)
 
     def __str__(self):
         return "Label(%d)" % self._count
@@ -55,6 +50,9 @@ def ufl_domains(self):
         "Return tuple of domains related to this terminal object."
         return ()
 
+    def _ufl_signature_data_(self, renumbering):
+        return ("Label", renumbering[self])
+
 
 @ufl_type(is_shaping=True, is_index_free=True, num_ops=1, inherit_shape_from_operand=0)
 class Variable(Operator):

From f3743afb3962733ada958909fe74acb1441bbf2e Mon Sep 17 00:00:00 2001
From: Connor Ward <c.ward20@imperial.ac.uk>
Date: Fri, 30 Jun 2023 11:17:16 +0100
Subject: [PATCH 2/3] fixup

---
 ufl/coefficient.py     |  2 +-
 ufl/constant.py        |  2 +-
 ufl/core/multiindex.py |  2 +-
 ufl/matrix.py          |  2 +-
 ufl/utils/counted.py   | 12 ++++++++++--
 ufl/variable.py        |  2 +-
 6 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/ufl/coefficient.py b/ufl/coefficient.py
index 61d2d2bad..5fec4e1f7 100644
--- a/ufl/coefficient.py
+++ b/ufl/coefficient.py
@@ -39,7 +39,7 @@ def __getnewargs__(self):
         return (self._ufl_function_space, self._count)
 
     def __init__(self, function_space, count=None):
-        Counted.__init__(self, count)
+        Counted.__init__(self, count, BaseCoefficient)
 
         if isinstance(function_space, FiniteElementBase):
             # For legacy support for .ufl files using cells, we map
diff --git a/ufl/constant.py b/ufl/constant.py
index 704604f7a..e5b612efe 100644
--- a/ufl/constant.py
+++ b/ufl/constant.py
@@ -20,7 +20,7 @@ class Constant(Terminal, Counted):
 
     def __init__(self, domain, shape=(), count=None):
         Terminal.__init__(self)
-        Counted.__init__(self, count)
+        Counted.__init__(self, count, Constant)
 
         self._ufl_domain = as_domain(domain)
         self._ufl_shape = shape
diff --git a/ufl/core/multiindex.py b/ufl/core/multiindex.py
index 5ae353703..ff10caacb 100644
--- a/ufl/core/multiindex.py
+++ b/ufl/core/multiindex.py
@@ -78,7 +78,7 @@ class Index(IndexBase, Counted):
 
     def __init__(self, count=None):
         IndexBase.__init__(self)
-        Counted.__init__(self, count)
+        Counted.__init__(self, count, Index)
 
     def __hash__(self):
         return hash(("Index", self._count))
diff --git a/ufl/matrix.py b/ufl/matrix.py
index 498b2fb5e..b5e15c653 100644
--- a/ufl/matrix.py
+++ b/ufl/matrix.py
@@ -37,7 +37,7 @@ def __getnewargs__(self):
 
     def __init__(self, row_space, column_space, count=None):
         BaseForm.__init__(self)
-        Counted.__init__(self, count)
+        Counted.__init__(self, count, Matrix)
 
         if not isinstance(row_space, AbstractFunctionSpace):
             raise ValueError("Expecting a FunctionSpace as the row space.")
diff --git a/ufl/utils/counted.py b/ufl/utils/counted.py
index 8528b30d6..0c7e306f0 100644
--- a/ufl/utils/counted.py
+++ b/ufl/utils/counted.py
@@ -19,9 +19,17 @@ class Counted:
 
     _counter = None
 
-    def __init__(self, count=None):
+    def __init__(self, count=None, cls=None):
+        """Initialize the Counted instance.
+
+        :arg count: The object count, if ``None`` defaults to the next value
+                    according to the global counter (per type).
+        :arg cls: Class to attach the global counter too. If ``None`` then
+                  ``type(self)`` will be used.
+
+        """
         # create a new counter for each subclass
-        cls = type(self)
+        cls = cls or type(self)
         if cls._counter is None:
             cls._counter = itertools.count()
 
diff --git a/ufl/variable.py b/ufl/variable.py
index 10e755b64..e009a946d 100644
--- a/ufl/variable.py
+++ b/ufl/variable.py
@@ -22,7 +22,7 @@ class Label(Terminal, Counted):
 
     def __init__(self, count=None):
         Terminal.__init__(self)
-        Counted.__init__(self, count)
+        Counted.__init__(self, count, Label)
 
     def __str__(self):
         return "Label(%d)" % self._count

From d54bc1821a36f07cc2ac1865cdd602cdc9ec6735 Mon Sep 17 00:00:00 2001
From: Connor Ward <c.ward20@imperial.ac.uk>
Date: Thu, 13 Jul 2023 10:32:57 +0100
Subject: [PATCH 3/3] fixup

---
 ufl/coefficient.py     |  3 ++-
 ufl/core/multiindex.py |  2 +-
 ufl/form.py            |  4 ++--
 ufl/matrix.py          |  1 +
 ufl/utils/counted.py   | 15 ++++++++-------
 ufl/variable.py        |  2 +-
 6 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/ufl/coefficient.py b/ufl/coefficient.py
index 5fec4e1f7..6440e15da 100644
--- a/ufl/coefficient.py
+++ b/ufl/coefficient.py
@@ -39,7 +39,7 @@ def __getnewargs__(self):
         return (self._ufl_function_space, self._count)
 
     def __init__(self, function_space, count=None):
-        Counted.__init__(self, count, BaseCoefficient)
+        Counted.__init__(self, count, Coefficient)
 
         if isinstance(function_space, FiniteElementBase):
             # For legacy support for .ufl files using cells, we map
@@ -108,6 +108,7 @@ class Cofunction(BaseCoefficient, BaseForm):
 
     __slots__ = (
         "_count",
+        "_counted_class",
         "_arguments",
         "_ufl_function_space",
         "ufl_operands",
diff --git a/ufl/core/multiindex.py b/ufl/core/multiindex.py
index ff10caacb..0847f2148 100644
--- a/ufl/core/multiindex.py
+++ b/ufl/core/multiindex.py
@@ -74,7 +74,7 @@ class Index(IndexBase, Counted):
     """UFL value: An index with no value assigned.
 
     Used to represent free indices in Einstein indexing notation."""
-    __slots__ = ("_count",)
+    __slots__ = ("_count", "_counted_class")
 
     def __init__(self, count=None):
         IndexBase.__init__(self)
diff --git a/ufl/form.py b/ufl/form.py
index 5124d9d13..63313f88a 100644
--- a/ufl/form.py
+++ b/ufl/form.py
@@ -449,10 +449,10 @@ def terminal_numbering(self):
         if self._terminal_numbering is None:
             exprs_by_type = defaultdict(set)
             for counted_expr in extract_type(self, Counted):
-                exprs_by_type[type(counted_expr)].add(counted_expr)
+                exprs_by_type[counted_expr._counted_class].add(counted_expr)
 
             numbering = {}
-            for type_, exprs in exprs_by_type.items():
+            for exprs in exprs_by_type.values():
                 for i, expr in enumerate(sorted_by_count(exprs)):
                     numbering[expr] = i
             self._terminal_numbering = numbering
diff --git a/ufl/matrix.py b/ufl/matrix.py
index b5e15c653..0976a1b0a 100644
--- a/ufl/matrix.py
+++ b/ufl/matrix.py
@@ -24,6 +24,7 @@ class Matrix(BaseForm, Counted):
 
     __slots__ = (
         "_count",
+        "_counted_class",
         "_ufl_function_spaces",
         "ufl_operands",
         "_repr",
diff --git a/ufl/utils/counted.py b/ufl/utils/counted.py
index 0c7e306f0..f04bc2aca 100644
--- a/ufl/utils/counted.py
+++ b/ufl/utils/counted.py
@@ -19,21 +19,22 @@ class Counted:
 
     _counter = None
 
-    def __init__(self, count=None, cls=None):
+    def __init__(self, count=None, counted_class=None):
         """Initialize the Counted instance.
 
         :arg count: The object count, if ``None`` defaults to the next value
                     according to the global counter (per type).
-        :arg cls: Class to attach the global counter too. If ``None`` then
-                  ``type(self)`` will be used.
+        :arg counted_class: Class to attach the global counter too. If ``None``
+                            then ``type(self)`` will be used.
 
         """
         # create a new counter for each subclass
-        cls = cls or type(self)
-        if cls._counter is None:
-            cls._counter = itertools.count()
+        counted_class = counted_class or type(self)
+        if counted_class._counter is None:
+            counted_class._counter = itertools.count()
 
-        self._count = count if count is not None else next(cls._counter)
+        self._count = count if count is not None else next(counted_class._counter)
+        self._counted_class = counted_class
 
     def count(self):
         return self._count
diff --git a/ufl/variable.py b/ufl/variable.py
index e009a946d..1338b2e85 100644
--- a/ufl/variable.py
+++ b/ufl/variable.py
@@ -18,7 +18,7 @@
 
 @ufl_type()
 class Label(Terminal, Counted):
-    __slots__ = ("_count",)
+    __slots__ = ("_count", "_counted_class")
 
     def __init__(self, count=None):
         Terminal.__init__(self)