From 884f40df11f1feb56d41f69fea773ffede768803 Mon Sep 17 00:00:00 2001
From: Krystine Sherwin <93062060+KrystalDelusion@users.noreply.github.com>
Date: Fri, 24 May 2024 03:42:12 +1200
Subject: [PATCH] write_btor: Include `$assert` and `$assume` cells in -ywmap
 output

---
 backends/btor/btor.cc | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/backends/btor/btor.cc b/backends/btor/btor.cc
index 9cfd967e5..35132fb2e 100644
--- a/backends/btor/btor.cc
+++ b/backends/btor/btor.cc
@@ -97,6 +97,8 @@ struct BtorWorker
 	vector<ywmap_btor_sig> ywmap_states;
 	dict<SigBit, int> ywmap_clock_bits;
 	dict<SigBit, int> ywmap_clock_inputs;
+	vector<Cell *> ywmap_asserts;
+	vector<Cell *> ywmap_assumes;
 
 
 	PrettyJson ywmap_json;
@@ -1277,6 +1279,8 @@ struct BtorWorker
 				btorf("%d or %d %d %d\n", nid_a_or_not_en, sid, nid_a, nid_not_en);
 				btorf("%d constraint %d\n", nid, nid_a_or_not_en);
 
+				if (ywmap_json.active()) ywmap_assumes.emplace_back(cell);
+
 				btorf_pop(log_id(cell));
 			}
 
@@ -1301,6 +1305,8 @@ struct BtorWorker
 					} else {
 						int nid = next_nid++;
 						btorf("%d bad %d%s\n", nid, nid_en_and_not_a, getinfo(cell, true).c_str());
+
+						if (ywmap_json.active()) ywmap_asserts.emplace_back(cell);
 					}
 				}
 
@@ -1458,6 +1464,7 @@ struct BtorWorker
 				log_assert(cursor == 0);
 				log_assert(GetSize(todo) == 1);
 				btorf("%d bad %d\n", nid, todo[cursor]);
+				// What do we do with ywmap_asserts when using single_bad?
 			}
 		}
 
@@ -1523,6 +1530,18 @@ struct BtorWorker
 				emit_ywmap_btor_sig(entry);
 			ywmap_json.end_array();
 
+			ywmap_json.name("asserts");
+			ywmap_json.begin_array();
+			for (Cell *cell : ywmap_asserts)
+				ywmap_json.value(witness_path(cell));
+			ywmap_json.end_array();
+
+			ywmap_json.name("assumes");
+			ywmap_json.begin_array();
+			for (Cell *cell : ywmap_assumes)
+				ywmap_json.value(witness_path(cell));
+			ywmap_json.end_array();
+
 			ywmap_json.end_object();
 		}
 	}