Moved smt2 yosys info parsing from smtbmc.py to smtio.py

3 files changed, 56 insertions, 16 deletions

diff --git a/backends/smt2/smt2.cc b/backends/smt2/smt2.cc
index 1ff1cdbf5..02d6f3fb6 100644
--- a/backends/smt2/smt2.cc
+++ b/backends/smt2/smt2.cc
@@ -932,6 +932,10 @@ struct Smt2Backend : public Backend {
 			worker.write(*f);
 		}
 
+		Module *topmod = design->top_module();
+		if (topmod)
+			*f << stringf("; yosys-smt2-topmod %s\n", log_id(topmod));
+
 		*f << stringf("; end of yosys output\n");
 
 		if (template_f.is_open()) {
diff --git a/backends/smt2/smtbmc.py b/backends/smt2/smtbmc.py
index 057776f6d..3d96b07a0 100644
--- a/backends/smt2/smtbmc.py
+++ b/backends/smt2/smtbmc.py
@@ -95,37 +95,28 @@ smt = smtio(opts=so)
 print("%s Solver: %s" % (smt.timestamp(), so.solver))
 smt.setup("QF_AUFBV")
 
-debug_nets = dict()
-debug_nets_re = re.compile(r"^; yosys-smt2-(input|output|register|wire) (\S+) (\d+)")
-current_module = None
-
 with open(args[0], "r") as f:
     for line in f:
-        match = debug_nets_re.match(line)
-        if match:
-            debug_nets[current_module].add(match.group(2))
-        if line.startswith("; yosys-smt2-module"):
-            current_module = line.split()[2]
-            debug_nets[current_module] = set()
         smt.write(line)
-    if topmod is None:
-        topmod = current_module
+        smt.getinfo(line)
 
-assert topmod is not None
-assert topmod in debug_nets
+if topmod is None:
+    topmod = smt.topmod
 
+assert topmod is not None
+assert topmod in smt.modinfo
 
 def write_vcd_model(steps):
     print("%s Writing model to VCD file." % smt.timestamp())
 
     vcd = mkvcd(open(vcdfile, "w"))
-    for netname in sorted(debug_nets[topmod]):
+    for netname in sorted(smt.modinfo[topmod].wsize.keys()):
         width = len(smt.get_net_bin(topmod, netname, "s0"))
         vcd.add_net(netname, width)
 
     for i in range(steps):
         vcd.set_time(i)
-        for netname in debug_nets[topmod]:
+        for netname in smt.modinfo[topmod].wsize.keys():
             vcd.set_net(netname, smt.get_net_bin(topmod, netname, "s%d" % i))
 
     vcd.set_time(steps)
diff --git a/backends/smt2/smtio.py b/backends/smt2/smtio.py
index 6e8bded77..14ad75e3e 100644
--- a/backends/smt2/smtio.py
+++ b/backends/smt2/smtio.py
@@ -22,6 +22,15 @@ import subprocess
 from select import select
 from time import time
 
+class smtmodinfo:
+    def __init__(self):
+        self.inputs = set()
+        self.outputs = set()
+        self.registers = set()
+        self.wires = set()
+        self.wsize = dict()
+        self.cells = dict()
+
 class smtio:
     def __init__(self, solver=None, debug_print=None, debug_file=None, timeinfo=None, opts=None):
         if opts is not None:
@@ -63,6 +72,10 @@ class smtio:
         self.p = subprocess.Popen(popen_vargs, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
         self.start_time = time()
 
+        self.modinfo = dict()
+        self.curmod = None
+        self.topmod = None
+
     def setup(self, logic="ALL", info=None):
         self.write("(set-logic %s)" % logic)
         if info is not None:
@@ -84,6 +97,38 @@ class smtio:
         self.p.stdin.write(bytes(stmt + "\n", "ascii"))
         self.p.stdin.flush()
 
+    def getinfo(self, stmt):
+        if not stmt.startswith("; yosys-smt2-"):
+            return
+
+        fields = stmt.split()
+
+        if fields[1] == "yosys-smt2-module":
+            self.curmod = fields[2]
+            self.modinfo[self.curmod] = smtmodinfo()
+
+        if fields[1] == "yosys-smt2-cell":
+            self.modinfo[self.curmod].cells[fields[3]] = fields[2]
+
+        if fields[1] == "yosys-smt2-topmod":
+            self.topmod = fields[2]
+
+        if fields[1] == "yosys-smt2-input":
+            self.modinfo[self.curmod].inputs.add(fields[2])
+            self.modinfo[self.curmod].wsize[fields[2]] = int(fields[3])
+
+        if fields[1] == "yosys-smt2-output":
+            self.modinfo[self.curmod].outputs.add(fields[2])
+            self.modinfo[self.curmod].wsize[fields[2]] = int(fields[3])
+
+        if fields[1] == "yosys-smt2-register":
+            self.modinfo[self.curmod].registers.add(fields[2])
+            self.modinfo[self.curmod].wsize[fields[2]] = int(fields[3])
+
+        if fields[1] == "yosys-smt2-wire":
+            self.modinfo[self.curmod].wires.add(fields[2])
+            self.modinfo[self.curmod].wsize[fields[2]] = int(fields[3])
+
     def read(self):
         stmt = []
         count_brackets = 0