From c8d0ab9afe8eab6fec0b22e80f263ab2e359bb95 Mon Sep 17 00:00:00 2001
From: Jonathan Coates <jonathan.coates@oxionics.com>
Date: Fri, 7 Feb 2025 00:33:08 +0000
Subject: [PATCH] Avoid quadratic behaviour in `_unify_attribute` (#2673)

ARTIQ maintains a list of all known instances of a particular class, and
after each attribute access, `_unify_attribute` is called to check all
these instances have this attribute (and it's the right type).

While the attribute type computation is cached, this is still O(nm)
(with n being the number of instances, and m the number of attribute
lookups). For something like ndscan's `FloatParamHandle.get`, you can
have a lot of both.

This commit changes `_unify_attribute` to only check the attributes of
new instances, effectively lowering the complexity to O(n). This
provides a small (5%) boost to compile times.

Signed-off-by: Jonathan Coates <jonathan.coates@oxionics.com>
---
 artiq/compiler/embedding.py | 33 ++++++++++++++++++++++++++++++---
 1 file changed, 30 insertions(+), 3 deletions(-)

diff --git a/artiq/compiler/embedding.py b/artiq/compiler/embedding.py
index 233d2b8fb..58ea63535 100644
--- a/artiq/compiler/embedding.py
+++ b/artiq/compiler/embedding.py
@@ -274,6 +274,27 @@ class EmbeddingMap:
             ))
 
 
+class _ValueInfo:
+    """
+    A collection of all values of a particular type.
+
+    Attributes:
+
+    :attr:`objects`: A list of all objects and the location they were added from.
+    :attr:`unchecked_attributes`: The known attributes for this type, and a list of
+        values where we have not yet checked the attribute's presence.
+    """
+    def __init__(self):
+        self.objects: list[tuple[object, source.Range]] = []
+        self.unchecked_attributes: dict[str, list[tuple[object, source.Range]]] = {}
+
+    def append(self, val_and_loc):
+        self.objects.append(val_and_loc)
+
+        for attr_store in self.unchecked_attributes.values():
+            attr_store.append(val_and_loc)
+
+
 class ASTSynthesizer:
     def __init__(self, embedding_map, value_map, quote_function=None, expanded_from=None):
         self.source = ""
@@ -807,7 +828,13 @@ class StitchingInferencer(Inferencer):
         # that we can successfully serialize the value of the attribute we
         # are now adding at the code generation stage.
         object_type = value_node.type.find()
-        for object_value, object_loc in self.value_map[object_type]:
+        values: _ValueInfo = self.value_map[object_type]
+
+        # Take all objects whose attribute we haven't checked yet.
+        attribute_objects = values.unchecked_attributes.get(attr_name, values.objects)
+        values.unchecked_attributes[attr_name] = []
+
+        for object_value, object_loc in attribute_objects:
             attr_type_key = (id(object_value), attr_name)
             try:
                 attributes, attr_value_type = self.attr_type_cache[attr_type_key]
@@ -888,7 +915,7 @@ class Stitcher:
         self.functions = {}
 
         self.embedding_map = EmbeddingMap(old_embedding_map)
-        self.value_map = defaultdict(lambda: [])
+        self.value_map = defaultdict(_ValueInfo)
         self.definitely_changed = False
 
         self.destination = destination
@@ -946,7 +973,7 @@ class Stitcher:
                     # value is guaranteed to have it too.
                     continue
 
-                for value, loc in self.value_map[instance_type]:
+                for value, loc in self.value_map[instance_type].objects:
                     if hasattr(value, attribute):
                         continue