From 94d7c22714dad09b24321835c36ec3100f73f9d7 Mon Sep 17 00:00:00 2001
From: Jannis Harder <me@jix.one>
Date: Thu, 14 Dec 2023 16:45:19 +0100
Subject: [PATCH] yosys-witness: Add aiw2yw --present-only to omit unused
 signals

---
 backends/smt2/witness.py | 53 +++++++++++++++++++++++++++++++---------
 1 file changed, 41 insertions(+), 12 deletions(-)

diff --git a/backends/smt2/witness.py b/backends/smt2/witness.py
index a39500c2d..b7e25851c 100644
--- a/backends/smt2/witness.py
+++ b/backends/smt2/witness.py
@@ -183,7 +183,8 @@ This requires a Yosys witness AIGER map file as generated by 'write_aiger -ywmap
 @click.argument("mapfile", type=click.File("r"))
 @click.argument("output", type=click.File("w"))
 @click.option("--skip-x", help="Leave input x bits unassigned.", is_flag=True)
-def aiw2yw(input, mapfile, output, skip_x):
+@click.option("--present-only", help="Only include bits present in at least one time step.", is_flag=True)
+def aiw2yw(input, mapfile, output, skip_x, present_only):
     input_name = input.name
     click.echo(f"Converting AIGER witness trace {input_name!r} to Yosys witness trace {output.name!r}...")
     click.echo(f"Using Yosys witness AIGER map file {mapfile.name!r}")
@@ -211,16 +212,23 @@ def aiw2yw(input, mapfile, output, skip_x):
     if not re.match(r'[0]*$', ffline):
         raise click.ClickException(f"{input_name}: non-default initial state not supported")
 
-    outyw = WriteWitness(output, 'yosys-witness aiw2yw')
+    if not present_only:
+        outyw = WriteWitness(output, "yosys-witness aiw2yw")
 
-    for clock in aiger_map.clocks:
-        outyw.add_clock(clock["path"], clock["offset"], clock["edge"])
+        for clock in aiger_map.clocks:
+            outyw.add_clock(clock["path"], clock["offset"], clock["edge"])
 
-    for (path, offset), id in aiger_map.sigmap.bit_to_id.items():
-        outyw.add_sig(path, offset, init_only=id in aiger_map.init_inputs)
+        for (path, offset), id in aiger_map.sigmap.bit_to_id.items():
+            outyw.add_sig(path, offset, init_only=id in aiger_map.init_inputs)
 
     missing = set()
+    seen = set()
 
+    buffered_steps = []
+
+    skip = "x?" if skip_x else "?"
+
+    t = -1
     while True:
         inline = next(input, None)
         if inline is None:
@@ -232,17 +240,20 @@ def aiw2yw(input, mapfile, output, skip_x):
         if inline.startswith("#"):
             continue
 
-        if not re.match(r'[01x]*$', ffline):
+        t += 1
+
+        if not re.match(r"[01x]*$", inline):
             raise click.ClickException(f"{input_name}: unexpected data in AIGER witness file")
 
         if len(inline) != aiger_map.input_count:
             raise click.ClickException(
-                    f"{input_name}: {mapfile.name}: number of inputs does not match, "
-                    f"{len(inline)} in witness, {aiger_map.input_count} in map file")
+                f"{input_name}: {mapfile.name}: number of inputs does not match, "
+                f"{len(inline)} in witness, {aiger_map.input_count} in map file"
+            )
 
         values = WitnessValues()
         for i, v in enumerate(inline):
-            if outyw.t > 0 and i in aiger_map.init_inputs:
+            if v in skip or (t > 0 and i in aiger_map.init_inputs):
                 continue
 
             try:
@@ -250,11 +261,29 @@ def aiw2yw(input, mapfile, output, skip_x):
             except IndexError:
                 bit = None
             if bit is None:
-                missing.insert(i)
+                missing.add(i)
+            elif present_only:
+                seen.add(i)
 
             values[bit] = v
 
-        outyw.step(values, skip_x=skip_x)
+        if present_only:
+            buffered_steps.append(values)
+        else:
+            outyw.step(values)
+
+    if present_only:
+        outyw = WriteWitness(output, "yosys-witness aiw2yw")
+
+        for clock in aiger_map.clocks:
+            outyw.add_clock(clock["path"], clock["offset"], clock["edge"])
+
+        for (path, offset), id in aiger_map.sigmap.bit_to_id.items():
+            if id in seen:
+                outyw.add_sig(path, offset, init_only=id in aiger_map.init_inputs)
+
+        for values in buffered_steps:
+            outyw.step(values)
 
     outyw.end_trace()