From 6469d90293b465201c9903f2dc1d02045cc33bcd Mon Sep 17 00:00:00 2001
From: Jannis Harder <me@jix.one>
Date: Wed, 7 Feb 2024 18:39:32 +0100
Subject: [PATCH] write_aiger: Include `$assert` and `$assume` cells in -ywmap
 output

---
 backends/aiger/aiger.cc | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/backends/aiger/aiger.cc b/backends/aiger/aiger.cc
index fe4f7681d..f2cb5d9bc 100644
--- a/backends/aiger/aiger.cc
+++ b/backends/aiger/aiger.cc
@@ -67,6 +67,8 @@ struct AigerWriter
 	int initstate_ff = 0;
 
 	dict<SigBit, int> ywmap_clocks;
+	vector<Cell *> ywmap_asserts;
+	vector<Cell *> ywmap_assumes;
 
 	int mkgate(int a0, int a1)
 	{
@@ -269,6 +271,7 @@ struct AigerWriter
 				unused_bits.erase(A);
 				unused_bits.erase(EN);
 				asserts.push_back(make_pair(A, EN));
+				ywmap_asserts.push_back(cell);
 				continue;
 			}
 
@@ -279,6 +282,7 @@ struct AigerWriter
 				unused_bits.erase(A);
 				unused_bits.erase(EN);
 				assumes.push_back(make_pair(A, EN));
+				ywmap_assumes.push_back(cell);
 				continue;
 			}
 
@@ -852,6 +856,19 @@ struct AigerWriter
 		for (auto &it : init_lines)
 			json.value(it.second);
 		json.end_array();
+
+		json.name("asserts");
+		json.begin_array();
+		for (Cell *cell : ywmap_asserts)
+			json.value(witness_path(cell));
+		json.end_array();
+
+		json.name("assumes");
+		json.begin_array();
+		for (Cell *cell : ywmap_assumes)
+			json.value(witness_path(cell));
+		json.end_array();
+
 		json.end_object();
 	}