smt2/smtbmc: Fix FF witness data for fine grained or multi chunk FFs

The witness metadata was missing fine grained FFs completely and for
coarse grained FFs where the output connection has multiple chunks it
lacked the offset of the chunk within the SMT expression. This fixes
both, the later by adding an "smtoffset" field to the metadata.

3 files changed, 22 insertions, 5 deletions

diff --git a/backends/smt2/smt2.cc b/backends/smt2/smt2.cc
index 54783cf1b..7434b13da 100644
--- a/backends/smt2/smt2.cc
+++ b/backends/smt2/smt2.cc
@@ -559,6 +559,9 @@ struct Smt2Worker
 		if (cell->type.in(ID($_FF_), ID($_DFF_P_), ID($_DFF_N_)))
 		{
 			registers.insert(cell);
+			SigBit q_bit = cell->getPort(ID::Q);
+			if (q_bit.is_wire())
+				decls.push_back(witness_signal("reg", 1, 0, "", idcounter, q_bit.wire));
 			makebits(stringf("%s#%d", get_id(module), idcounter), 0, log_signal(cell->getPort(ID::Q)));
 			register_bool(cell->getPort(ID::Q), idcounter++);
 			recursive_cells.erase(cell);
@@ -589,9 +592,12 @@ struct Smt2Worker
 			if (cell->type.in(ID($ff), ID($dff)))
 			{
 				registers.insert(cell);
-				for (auto chunk : cell->getPort(ID::Q).chunks())
+				int smtoffset = 0;
+				for (auto chunk : cell->getPort(ID::Q).chunks()) {
 					if (chunk.is_wire())
-						decls.push_back(witness_signal("reg", chunk.width, chunk.offset, "", idcounter, chunk.wire));
+						decls.push_back(witness_signal("reg", chunk.width, chunk.offset, "", idcounter, chunk.wire, smtoffset));
+					smtoffset += chunk.width;
+				}
 				makebits(stringf("%s#%d", get_id(module), idcounter), GetSize(cell->getPort(ID::Q)), log_signal(cell->getPort(ID::Q)));
 				register_bv(cell->getPort(ID::Q), idcounter++);
 				recursive_cells.erase(cell);
@@ -1490,7 +1496,7 @@ struct Smt2Worker
 		return path;
 	}
 
-	std::string witness_signal(const char *type, int width, int offset, const std::string &smtname, int smtid, RTLIL::Wire *wire)
+	std::string witness_signal(const char *type, int width, int offset, const std::string &smtname, int smtid, RTLIL::Wire *wire, int smtoffset = 0)
 	{
 		std::vector<std::string> hiername;
 		const char *wire_name = wire->name.c_str();
@@ -1508,6 +1514,7 @@ struct Smt2Worker
 			{ "offset", offset },
 			{ "width", width },
 			{ "smtname", smtname.empty() ? json11::Json(smtid) : json11::Json(smtname) },
+			{ "smtoffset", smtoffset },
 			{ "path", witness_path(wire) },
 		}}).dump(line);
 		line += "\n";
diff --git a/backends/smt2/smtbmc.py b/backends/smt2/smtbmc.py
index 5f05287de..ac329053f 100644
--- a/backends/smt2/smtbmc.py
+++ b/backends/smt2/smtbmc.py
@@ -669,7 +669,7 @@ if inywfile is not None:
                         if common_end <= common_offset:
                             continue
 
-                        smt_expr = smt.net_expr(topmod, f"s{t}", wire["smtpath"])
+                        smt_expr = smt.witness_net_expr(topmod, f"s{t}", wire)
 
                         if not smt_bool:
                             slice_high = common_end - offset - 1
@@ -1267,7 +1267,8 @@ def write_yw_trace(steps_start, steps_stop, index, allregs=False):
                 sigs = seqs
 
             for sig in sigs:
-                step_values[sig["sig"]] = smt.bv2bin(smt.get(smt.net_expr(topmod, f"s{k}", sig["smtpath"])))
+                value = smt.bv2bin(smt.get(smt.witness_net_expr(topmod, f"s{k}", sig)))
+                step_values[sig["sig"]] = value
             yw.step(step_values)
 
         yw.end_trace()
diff --git a/backends/smt2/smtio.py b/backends/smt2/smtio.py
index 9af454cca..f5cfe436d 100644
--- a/backends/smt2/smtio.py
+++ b/backends/smt2/smtio.py
@@ -958,6 +958,15 @@ class SmtIo:
         nextbase = "(|%s_h %s| %s)" % (mod, path[0], base)
         return self.net_expr(nextmod, nextbase, path[1:])
 
+    def witness_net_expr(self, mod, base, witness):
+        net = self.net_expr(mod, base, witness["smtpath"])
+        is_bool = self.net_width(mod, witness["smtpath"]) == 1
+        if is_bool:
+            assert witness["width"] == 1
+            assert witness["smtoffset"] == 0
+            return net
+        return "((_ extract %d %d) %s)" % (witness["smtoffset"] + witness["width"] - 1, witness["smtoffset"], net)
+
     def net_width(self, mod, net_path):
         for i in range(len(net_path)-1):
             assert mod in self.modinfo