From 9c338a3ae42922c5dfbd5d4b18aca4a40463baa4 Mon Sep 17 00:00:00 2001
From: Pavol Rusnak <pavol@rusnak.io>
Date: Tue, 16 Jun 2020 10:36:50 +0000
Subject: [PATCH] core: add memprofile to profiler

---
 core/plot.py      | 13 +++++++++++++
 core/prof/prof.py | 38 ++++++++++++++++++++++++++++----------
 2 files changed, 41 insertions(+), 10 deletions(-)
 create mode 100644 core/plot.py

diff --git a/core/plot.py b/core/plot.py
new file mode 100644
index 000000000..92426a5f3
--- /dev/null
+++ b/core/plot.py
@@ -0,0 +1,13 @@
+import pandas as pd
+import plotly.express as px
+import plotly.graph_objects as go
+
+df = pd.read_csv("src/.memprofile", sep=";")
+
+fig = go.Figure()
+
+fig.add_trace(go.Scatter(x=df.trace, y=df.mem_total, name="mem_total"))
+fig.add_trace(go.Scatter(x=df.trace, y=df.mem_current, name="mem_current"))
+fig.add_trace(go.Scatter(x=df.trace, y=df.mem_peak, name="mem_peak"))
+
+fig.show()
diff --git a/core/prof/prof.py b/core/prof/prof.py
index 75703ed05..192195f2e 100644
--- a/core/prof/prof.py
+++ b/core/prof/prof.py
@@ -1,8 +1,15 @@
+ENABLE_COVERAGE = True
+ENABLE_MEMPROFILE = False
+
 import sys
 
 from uio import open
 from uos import getenv
 
+if ENABLE_MEMPROFILE:
+    from micropython import mem_total, mem_current, mem_peak
+
+
 # We need to insert "" to sys.path so that the frozen build can import main from the
 # frozen modules, and regular build can import it from current directory.
 sys.path.insert(0, "")
@@ -33,16 +40,24 @@ class Coverage:
 class _Prof:
     trace_count = 0
     display_flags = 0
-    __coverage = Coverage()
+    if ENABLE_COVERAGE:
+        __coverage = Coverage()
+    if ENABLE_MEMPROFILE:
+        __memprofile = open(".memprofile", "wt")
+        __memprofile.write("trace;mem_total;mem_current;mem_peak\n")
 
     def trace_tick(self, frame, event, arg):
         self.trace_count += 1
 
-        # if frame.f_code.co_filename.endswith('/loop.py'):
-        #     print(event, frame.f_code.co_filename, frame.f_lineno)
+        if ENABLE_COVERAGE:
+            if event == "line":
+                self.__coverage.line_tick(frame.f_code.co_filename, frame.f_lineno)
 
-        if event == "line":
-            self.__coverage.line_tick(frame.f_code.co_filename, frame.f_lineno)
+        if ENABLE_MEMPROFILE:
+            self.__memprofile.write(
+                "%d;%d;%d;%d\n"
+                % (self.trace_count, mem_total(), mem_current(), mem_peak())
+            )
 
     def coverage_data(self):
         return self.__coverage.lines_execution()
@@ -56,11 +71,14 @@ def trace_handler(frame, event, arg):
 def atexit():
     print("\n------------------ script exited ------------------")
     print("Total traces executed: ", __prof__.trace_count)
-    with open(".coverage", "w") as f:
-        # wtf so private much beautiful wow
-        f.write("!coverage.py: This is a private format, don't read it directly!")
-        # poormans json
-        f.write(str(__prof__.coverage_data()).replace("'", '"'))
+    if ENABLE_COVERAGE:
+        with open(".coverage", "w") as f:
+            # wtf so private much beautiful wow
+            f.write("!coverage.py: This is a private format, don't read it directly!")
+            # poormans json
+            f.write(str(__prof__.coverage_data()).replace("'", '"'))
+    if ENABLE_MEMPROFILE:
+        __prof__.__memprofile.close()
 
 
 sys.atexit(atexit)