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.

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";

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()

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