From d6a37c466061be4d51b2fd60ca21c480dbe2a4dd Mon Sep 17 00:00:00 2001
From: Andreas Sandberg <andreas.sandberg@arm.com>
Date: Fri, 21 Aug 2020 15:55:53 +0100
Subject: [PATCH] python: Add support for introspecting scalar stats

This change adds a wrapper for the ScalarInfo stat type to enable
introspection of scalar stats from Python. Due to the slightly
confusing use of proxy objects in the stat system, PyBind11 fails to
automatically cast to the right wrapper type. This is worked around in
the by explicitly casting to the relevant type's Python wrapper.

To make the interface more Python-friendly, this change also changes
the semantics of resolveStat to raise an exception if the stat can't
be found.

Change-Id: If1fc6fe238fc9d69d4e22369a4988a06407d2f7c
Signed-off-by: Andreas Sandberg <andreas.sandberg@arm.com>
Reviewed-on: https://gem5-review.googlesource.com/c/public/gem5/+/33176
Reviewed-by: Jason Lowe-Power <power.jg@gmail.com>
Maintainer: Jason Lowe-Power <power.jg@gmail.com>
Tested-by: kokoro <noreply+kokoro@google.com>
---
 src/python/pybind11/stats.cc | 48 ++++++++++++++++++++++++++++++++++--
 1 file changed, 46 insertions(+), 2 deletions(-)

diff --git a/src/python/pybind11/stats.cc b/src/python/pybind11/stats.cc
index 1149eba3a..b146aa306 100644
--- a/src/python/pybind11/stats.cc
+++ b/src/python/pybind11/stats.cc
@@ -54,6 +54,25 @@
 
 namespace py = pybind11;
 
+static const py::object
+cast_stat_info(const Stats::Info *info)
+{
+    /* PyBind11 gets confused by the InfoProxy magic, so we need to
+     * explicitly cast to the right wrapper type. */
+
+#define TRY_CAST(T) do {                                \
+        auto _stat = dynamic_cast<const T *>(info);     \
+        if (_stat)                                      \
+            return py::cast(_stat);                     \
+    } while (0)
+
+    TRY_CAST(Stats::ScalarInfo);
+
+    return py::cast(info);
+
+#undef TRY_CAST
+}
+
 namespace Stats {
 
 void
@@ -120,14 +139,39 @@ pybind_init_stats(py::module &m_native)
         .def("visit", &Stats::Info::visit)
         ;
 
+    py::class_<Stats::ScalarInfo, Stats::Info,
+               std::unique_ptr<Stats::ScalarInfo, py::nodelete>>(
+                   m, "ScalarInfo")
+        .def("value", &Stats::ScalarInfo::value)
+        .def("result", &Stats::ScalarInfo::result)
+        .def("total", &Stats::ScalarInfo::total)
+        ;
+
     py::class_<Stats::Group, std::unique_ptr<Stats::Group, py::nodelete>>(
         m, "Group")
         .def("regStats", &Stats::Group::regStats)
         .def("resetStats", &Stats::Group::resetStats)
         .def("preDumpStats", &Stats::Group::preDumpStats)
-        .def("getStats", &Stats::Group::getStats)
+        .def("getStats", [](const Stats::Group &self)
+             -> std::vector<py::object> {
+
+                 auto stats = self.getStats();
+                std::vector<py::object> py_stats;
+                py_stats.reserve(stats.size());
+                std::transform(stats.begin(), stats.end(),
+                               std::back_inserter(py_stats),
+                               cast_stat_info);
+                return py_stats;
+            })
         .def("getStatGroups", &Stats::Group::getStatGroups)
         .def("addStatGroup", &Stats::Group::addStatGroup)
-        .def("resolveStat", &Stats::Group::resolveStat)
+        .def("resolveStat", [](const Stats::Group &self,
+                               const std::string &name) -> py::object {
+                 const Stats::Info *stat = self.resolveStat(name);
+                 if (!stat)
+                     throw pybind11::key_error("Unknown stat name");
+
+                 return cast_stat_info(stat);
+             })
         ;
 }
-- 
2.30.2