refactor: review feedback

src/puya/ir/builder/main.py

@@ -93,7 +93,7 @@ def __init__(
         self.context = context.for_function(function, subroutine, self)
         self._itxn = InnerTransactionBuilder(self.context)
         self._single_eval_cache = dict[awst_nodes.SingleEvaluation, TExpression]()
-        self._visited_exprs = dict[_IdentityEquality, TExpression]()
+        self._visited_exprs = dict[tuple[int, awst_nodes.Expression], TExpression]()
 
     @classmethod
     def build_body(
@@ -1197,7 +1197,9 @@ def visit_and_materialise_single(
     def visit_and_materialise(
         self, expr: awst_nodes.Expression, temp_description: str | Sequence[str] = "tmp"
     ) -> Sequence[Value]:
-        value_seq_or_provider = self._visit_and_check_for_double_eval(expr, temp_description)
+        value_seq_or_provider = self._visit_and_check_for_double_eval(
+            expr, materialise_as=temp_description
+        )
         if value_seq_or_provider is None:
             raise InternalError(
                 "No value produced by expression IR conversion", expr.source_location
@@ -1214,16 +1216,16 @@ def visit_expr(self, expr: awst_nodes.Expression) -> ValueProvider:
         return value_seq_or_provider
 
     def _visit_and_check_for_double_eval(
-        self, expr: awst_nodes.Expression, desc: str | Sequence[str] | None = None
+        self, expr: awst_nodes.Expression, *, materialise_as: str | Sequence[str] | None = None
     ) -> ValueProvider | None:
         # explicit SingleEvaluation nodes already handle this
         if isinstance(expr, awst_nodes.SingleEvaluation):
             return expr.accept(self)
-        # we need a wrapper object to provide identity equality for the expression, and
-        # to ensure the lifetime of the expression is as long as the cache.
+        # include the expression in the key to ensure the lifetime of the
+        # expression is as long as the cache.
         # Temporary nodes may end up with the same id if nothing is referencing them
         # e.g. such as used in _update_implicit_out_var
-        expr_id = _IdentityEquality(expr)
+        expr_id = id(expr), expr
         try:
             result = self._visited_exprs[expr_id]
         except KeyError:
@@ -1241,10 +1243,10 @@ def _visit_and_check_for_double_eval(
             )
             return result
         source = expr.accept(self)
-        if desc is None or source is None or not source.types or isinstance(source, Value):
+        if materialise_as is None or not (source and source.types):
             result = source
         else:
-            values = self.materialise_value_provider(source, description=desc)
+            values = self.materialise_value_provider(source, description=materialise_as)
             if len(values) == 1:
                 (result,) = values
             else:
@@ -1382,17 +1384,3 @@ def get_comparison_op_for_wtype(
     raise InternalError(
         f"unsupported operation of {numeric_comparison_equivalent} on type of {wtype}"
     )
-
-
-class _IdentityEquality:
-
-    def __init__(self, obj: object):
-        self.obj = obj
-
-    def __hash__(self) -> int:
-        return id(self.obj)
-
-    def __eq__(self, other: object) -> bool:
-        if not isinstance(other, _IdentityEquality):
-            return NotImplemented
-        return self.obj is other.obj