From 05ab45f39b8a39af2d4dbb1b62393a54b7e268b3 Mon Sep 17 00:00:00 2001
From: Bugra Onal <onalbugra@gmail.com>
Date: Tue, 30 Aug 2022 09:15:35 -0700
Subject: [PATCH] Added graph store and read functionality

---
 compiler/base/timing_graph.py | 30 ++++++++++++++++++++++++++++++
 compiler/memchar.py           |  3 +++
 compiler/modules/sram.py      | 11 +++++++++++
 compiler/options.py           |  2 ++
 4 files changed, 46 insertions(+)

diff --git a/compiler/base/timing_graph.py b/compiler/base/timing_graph.py
index 46d7b518..9c6301f5 100644
--- a/compiler/base/timing_graph.py
+++ b/compiler/base/timing_graph.py
@@ -1,6 +1,7 @@
 import copy
 from collections import defaultdict
 import debug
+import json
 
 
 class timing_graph():
@@ -139,6 +140,35 @@ class timing_graph():
 
         return [self.edge_mods[(path[i], path[i+1])] for i in range(len(path)-1)]
 
+    def write(self, filename):
+        """
+        Export graph to a JSON file
+        """
+        # TODO: Find a proper way to store edge_mods values
+        with open(filename, 'w') as f:
+            f.write(
+                json.dumps(
+                    {
+                        'graph':
+                            {key: list(val) for key, val in self.graph.items()},
+                        'edge_mods':
+                            {
+                                ', '.join(key): str(value)
+                                for key, value in self.edge_mods.items()
+                            }
+                    }
+                )
+            )
+
+    def read(self, filename):
+        """
+        Read graph from JSON file
+        """
+        with open(filename, 'r') as f:
+            d = json.loads(f.read())
+            self.graph = {key: set(value) for key, value in d['graph'].items()}
+            self.edge_mods = {tuple(key.split(', ')): value for key, value in d['edge_mods'].items()}
+
     def __str__(self):
         """ override print function output """
 
diff --git a/compiler/memchar.py b/compiler/memchar.py
index c4ce6dad..425a8abd 100755
--- a/compiler/memchar.py
+++ b/compiler/memchar.py
@@ -56,6 +56,9 @@ s.setup_multiport_constants()
 OPTS.netlist_only = True
 OPTS.check_lvsdrc = False
 
+# TODO: remove this after adding trimmed netlist gen to sram run
+OPTS.trim_netlist = False
+
 # Characterize the design
 start_time = datetime.datetime.now()
 from characterizer import lib
diff --git a/compiler/modules/sram.py b/compiler/modules/sram.py
index c6f73295..2a722709 100644
--- a/compiler/modules/sram.py
+++ b/compiler/modules/sram.py
@@ -9,6 +9,7 @@ import datetime
 import os
 import debug
 from characterizer import functional
+from base import timing_graph
 from globals import OPTS, print_time
 import shutil
 
@@ -187,3 +188,13 @@ class sram():
             debug.print_raw("Extended Config: Writing to {0}".format(oname))
             self.extended_config_write(oname)
             print_time("Extended Config", datetime.datetime.now(), start_time)
+
+        # Write the graph if specified
+        if OPTS.write_graph:
+            start_time = datetime.datetime.now()
+            oname = OPTS.output_path + OPTS.output_name + "_graph.json"
+            debug.print_raw("Graph: Writing to {0}".format(oname))
+            graph = timing_graph()
+            self.s.build_graph(graph, self.name, self.s.pins)
+            graph.write(oname)
+            print_time("Graph", datetime.datetime.now(), start_time)
diff --git a/compiler/options.py b/compiler/options.py
index 516f9817..6c518d09 100644
--- a/compiler/options.py
+++ b/compiler/options.py
@@ -119,6 +119,8 @@ class options(optparse.Values):
     # Determines which analytical model to use.
     # Available Models: elmore, linear_regression
     model_name = "elmore"
+    # Write graph to a file
+    write_graph = False
 
     ###################
     # Tool options