wibox.hierarchy: Add ability to count widgets

This adds new code so that we can count how often a specific widget is
visible inside of all widget hierarchies.

Signed-off-by: Uli Schlachter <psychon@znc.in>

lib/wibox/hierarchy.lua

@@ -16,6 +16,16 @@ local no_parent = base.no_parent_I_know_what_I_am_doing
 
 local hierarchy = {}
 
+local widgets_to_count = setmetatable({}, { __mode = "k" })
+
+--- Add a widget to the list of widgets for which hierarchies should count their
+-- occurrences. Note that for correct operations, the widget must not yet be
+-- visible in any hierarchy.
+-- @param widget The widget that should be counted.
+function hierarchy.count_widget(widget)
+    widgets_to_count[widget] = true
+end
+
 local function hierarchy_new(redraw_callback, layout_callback, callback_arg)
     local result = {
         _matrix = matrix.identity,
@@ -37,7 +47,8 @@ local function hierarchy_new(redraw_callback, layout_callback, callback_arg)
             height = 0
         },
         _parent = nil,
-        _children = {}
+        _children = {},
+        _widget_counts = {},
     }
 
     function result._redraw()
@@ -144,6 +155,17 @@ function hierarchy_update(self, context, widget, width, height, region, matrix_t
         height = y2 - y1
     }
 
+    -- Update widget counts
+    self._widget_counts = {}
+    if widgets_to_count[self._widget] then
+        self._widget_counts[self._widget] = 1
+    end
+    for _, h in ipairs(self._children) do
+        for w, count in pairs(h._widget_counts) do
+            self._widget_counts[w] = (self._widget_counts[w] or 0) + count
+        end
+    end
+
     -- Check which part needs to be redrawn
 
     -- Are there any children which were removed? Their area needs a redraw.
@@ -260,6 +282,14 @@ function hierarchy:get_children()
     return self._children
 end
 
+--- Count how often this widget is visible inside this hierarchy. This function
+-- only works with widgets registered via `count_widget`.
+-- @param widget The widget that should be counted
+-- @return The number of times that this widget is contained in this hierarchy.
+function hierarchy:get_count(widget)
+    return self._widget_counts[widget] or 0
+end
+
 --- Does the given cairo context have an empty clip (aka "no drawing possible")?
 local function empty_clip(cr)
     local _, _, width, height = cr:clip_extents()

spec/wibox/hierarchy_spec.lua

@@ -259,6 +259,66 @@ describe("wibox.hierarchy", function()
             assert.is.same({ rect.x, rect.y, rect.width, rect.height }, { 4, 0, 5, 2 })
         end)
     end)
+
+    describe("widget counts", function()
+        local child, intermediate, parent
+        local unrelated
+        local context, instance
+        before_each(function()
+            local function nop() end
+            context = {}
+            child = make_widget(nil)
+            intermediate = make_widget({
+                make_child(child, 10, 20, matrix.identity)
+            })
+            parent = make_widget({
+                make_child(intermediate, 10, 20, matrix.identity)
+            })
+            unrelated = make_widget(nil)
+
+            hierarchy.count_widget(child)
+            hierarchy.count_widget(parent)
+            hierarchy.count_widget(unrelated)
+            instance = hierarchy.new(context, parent, 10, 20, nop, nop)
+        end)
+
+        it("basic counts", function()
+            local unrelated_other = make_widget(nil)
+            assert.is.equal(1, instance:get_count(child))
+            -- intermediate was not passed to hierarchy.count_widget()!
+            assert.is.equal(0, instance:get_count(intermediate))
+            assert.is.equal(1, instance:get_count(parent))
+            assert.is.equal(0, instance:get_count(unrelated))
+            assert.is.equal(0, instance:get_count(unrelated_other))
+        end)
+
+        it("after update", function()
+            -- Replace child and intermediate by just a new_child
+            local new_child = make_widget(nil)
+            parent.layout = function()
+                return { make_child(new_child, 10, 20, matrix.identity) }
+            end
+            parent:emit_signal("widget::layout_changed")
+            instance:update(context, parent, 10, 20)
+
+            assert.is.equal(0, instance:get_count(child))
+            -- new_child was not passed to hierarchy.count_widget()!
+            assert.is.equal(0, instance:get_count(new_child))
+            assert.is.equal(1, instance:get_count(parent))
+            assert.is.equal(0, instance:get_count(unrelated))
+        end)
+
+        it("collectible", function()
+            -- This test that hierarchy.count_widget() does not prevent garbage collection of the widget.
+            local weak = setmetatable({}, { __mode = "v"})
+            weak[1], weak[2], weak[3], weak[4] = child, intermediate, parent, instance
+            child, intermediate, parent, instance = nil, nil, nil, nil
+
+            assert.is.equal(4, #weak)
+            collectgarbage("collect")
+            assert.is.equal(0, #weak)
+        end)
+    end)
 end)
 
 -- vim: filetype=lua:expandtab:shiftwidth=4:tabstop=8:softtabstop=4:textwidth=80

spec/wibox/test_utils.lua

@@ -96,6 +96,12 @@ return {
             end
         end
         w._private.widget_caches = {}
+        w:connect_signal("widget::layout_changed", function()
+            -- TODO: This is not completely correct, since our parent's caches
+            -- are not cleared. For the time being, tests just have to handle
+            -- this clearing-part themselves.
+            w._private.widget_caches = {}
+        end)
 
         return w
     end,