From 4f04ad2acd3a1b28513790a378c1f7675fbab2fc Mon Sep 17 00:00:00 2001
From: Joseph Schorr <joseph.schorr@coreos.com>
Date: Wed, 4 Mar 2015 16:53:22 -0500
Subject: [PATCH] Change ImageTree to only use a single loop over the images
 when building. This should be slightly faster on large image sets

---
 util/imagetree.py | 31 +++++++++++++++----------------
 1 file changed, 15 insertions(+), 16 deletions(-)

diff --git a/util/imagetree.py b/util/imagetree.py
index 39cd5c3c9..a45ac29d1 100644
--- a/util/imagetree.py
+++ b/util/imagetree.py
@@ -1,14 +1,15 @@
 class ImageTreeNode(object):
   """ A node in the image tree. """
-  def __init__(self, image):
+  def __init__(self, image, child_map):
     self.image = image
     self.parent = None
-    self.children = []
     self.tags = []
 
-  def add_child(self, child):
-    self.children.append(child)
-    child.parent = self
+    self._child_map = child_map
+
+  @property
+  def children(self):
+    return self._child_map.get(str(self.image.id), [])
 
   def add_tag(self, tag):
     self.tags.append(tag)
@@ -18,8 +19,8 @@ class ImageTree(object):
   """ In-memory tree for easy traversal and lookup of images in a repository. """
 
   def __init__(self, all_images, all_tags, base_filter=None):
-    self._tag_map = {}
     self._image_map = {}
+    self._child_map = {}
 
     self._build(all_images, all_tags, base_filter)
 
@@ -33,18 +34,17 @@ class ImageTree(object):
         if image.id != base_filter and not str(base_filter) in ancestors:
           continue
 
-      self._image_map[image.id] = ImageTreeNode(image)
+      # Create the node for the image.
+      image_node = ImageTreeNode(image, self._child_map)
+      self._image_map[image.id] = image_node
 
-    # Connect the nodes to their parents.
-    for image_node in self._image_map.values():
-      image = image_node.image
+      # Add the node to the child map for its parent image (if any).
       parent_image_id = image.ancestors.split('/')[-2] if image.ancestors else None
-      if not parent_image_id:
-        continue
+      if parent_image_id:
+        if not parent_image_id in self._child_map:
+          self._child_map[parent_image_id] = []
 
-      parent_node = self._image_map.get(int(parent_image_id))
-      if parent_node is not None:
-        parent_node.add_child(image_node)
+        self._child_map[parent_image_id].append(image_node)
 
     # Build the tag map.
     for tag in all_tags:
@@ -52,7 +52,6 @@ class ImageTree(object):
       if not image_node:
         continue
 
-      self._tag_map = image_node
       image_node.add_tag(tag.name)