From ef9e672eb9d6bbc4aee7bb3c5c8bad03c5d796a2 Mon Sep 17 00:00:00 2001
From: Jason Lowe-Power <jason@lowepower.com>
Date: Mon, 1 Mar 2021 14:51:15 -0800
Subject: [PATCH] python: Add search functions to pystats groups

This change adds three functions, a `children` function which will
iterate through all of the children of group based (optionally) on some
predicate. Then, it implements a `find` function and a `find_re`
function using the `children` function.

The `find` function allows users to match statistics or groups
within a group. For instance, you might want to find all of the groups
within the system which have the name "cpu{i}". This is useful for
aggregate statistic values across multiple components.

Example:
total_instruuctions = sum([cpu.exec_context.thread_0.numInsts.value
                           for cpu in simstat.system.find('cpu')])

The find function matches based on substring. If the name given the find
function is a substring of the stat name or the group name the
stat/group will be returned.

The `find_re` function is the same as find, but matches a regular
expression instead of a simple substring match.

Note: this was originally reviewed on
https://gem5-review.googlesource.com/c/public/gem5/+/41603 was rebased
incorrectly before merging. This change fixes the rebase and adds back
the children() and re_find() functions.

Change-Id: Idaa1e9efc56fd26de3285d3fa505087ddd78ac8a
Signed-off-by: Jason Lowe-Power <jason@lowepower.com>
Reviewed-on: https://gem5-review.googlesource.com/c/public/gem5/+/42014
Maintainer: Jason Lowe-Power <power.jg@gmail.com>
Tested-by: kokoro <noreply+kokoro@google.com>
Reviewed-by: Andreas Sandberg <andreas.sandberg@arm.com>
---
 src/python/m5/ext/pystats/group.py | 68 +++++++++++++++++++++++++-----
 1 file changed, 58 insertions(+), 10 deletions(-)

diff --git a/src/python/m5/ext/pystats/group.py b/src/python/m5/ext/pystats/group.py
index 10887e20a1..bde1c40ca2 100644
--- a/src/python/m5/ext/pystats/group.py
+++ b/src/python/m5/ext/pystats/group.py
@@ -24,7 +24,8 @@
 # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-from typing import Dict, Iterator, List, Optional, Union
+import re
+from typing import Callable, Dict, Iterator, List, Optional, Union
 
 from .jsonserializable import JsonSerializable
 from .statistic import Scalar, Statistic
@@ -53,25 +54,72 @@ class Group(JsonSerializable):
         for key,value in kwargs.items():
             setattr(self, key, value)
 
+    def children(self, predicate: Optional[Callable[[str], bool]] = None
+                 ) -> Iterator[Union["Group", Statistic]]:
+        """ Iterate through all of the children, optionally with a predicate
+
+        ```
+        >>> system.children(lambda _name: 'cpu' in name)
+        [cpu0, cpu1, cpu2]
+        ```
+
+        :param: predicate(str) -> bool: Optional. Each child's name is passed
+                to this function. If it returns true, then the child is
+                yielded. Otherwise, the child is skipped.
+                If not provided then all children are returned.
+        """
+        for attr in self.__dict__:
+            # Check the provided predicate. If not a match, skip this child
+            if predicate and not predicate(attr): continue
+            obj = getattr(self, attr)
+            if isinstance(obj, Group) or isinstance(obj, Statistic):
+                yield obj
+
     def find(self, name: str) -> Iterator[Union["Group", Statistic]]:
         """ Find all stats that match the name
+
         This function searches all of the "children" in this group. It yields
         the set of attributes (children) that have the `name` as a substring.
         The order of the objects returned by the generator is arbitrary.
+
         ```
-        system.find('cpu') -> [cpu0, cpu1, cpu2, cpu3, other_cpu, ...]
+        >>> system.find('cpu')
+        [cpu0, cpu1, cpu2, cpu3, other_cpu, ...]
         ```
+
         This is useful for performing aggregates over substats. For instance:
+
         ```
-        total_instruuctions = sum([cpu.exec_context.thread_0.numInsts.value
-                                   for cpu in simstat.system.find('cpu')])
+        >>> total_instructions = sum([cpu.exec_context.thread_0.numInsts.value
+                                      for cpu in simstat.system.find('cpu')])
+        100000
         ```
+
+        :param: name: The name to search for
         """
-        for attr in self.__dict__:
-            if name in attr:
-                obj = getattr(self, attr)
-                if isinstance(obj, Group) or isinstance(obj, Statistic):
-                    yield obj
+        yield from self.children(lambda _name: _name in name)
+
+    def find_re(self, regex: Union[str, re.Pattern]
+                ) -> Iterator[Union["Group", Statistic]]:
+        """ Find all stats that match the name
+
+        This function searches all of the "children" in this group. It yields
+        the set of attributes (children) that have the `name` mathing the
+        regex provided. The order of the objects returned by the generator is
+        arbitrary.
+
+        ```
+        >>> system.find_re('cpu[0-9]')
+        [cpu0, cpu1, cpu2]
+        ```
+        Note: The above will not match `cpu_other`.
+
+        :param: regex: The regular expression used to search. Can be a
+                precompiled regex or a string in regex format
+        """
+        if isinstance(regex, str):
+            regex = re.compile(regex)
+        yield from self.children(lambda _name: regex.search(_name))
 
 class Vector(Group):
     """
@@ -86,4 +134,4 @@ class Vector(Group):
                                      type="Vector",
                                      time_conversion=None,
                                      **scalar_map,
-                                    )
\ No newline at end of file
+                                    )