1 files changed, 148 insertions, 2 deletions
diff --git a/backends/smt2/witness.py b/backends/smt2/witness.py
index 1105d9ed0..03d72a17b 100644
--- a/backends/smt2/witness.py
+++ b/backends/smt2/witness.py
@@ -17,11 +17,12 @@
 # OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 #
 
-import os, sys
+import os, sys, itertools, re
 ##yosys-sys-path##
+import json
 import click
 
-from ywio import ReadWitness, WriteWitness, WitnessSig
+from ywio import ReadWitness, WriteWitness, WitnessSig, WitnessSigMap, WitnessValues, coalesce_signals
 
 @click.group()
 def cli():
@@ -86,5 +87,150 @@ def yw2yw(input, output):
 
     click.echo(f"Copied {outyw.t + 1} time steps.")
 
+
+class AigerMap:
+    def __init__(self, mapfile):
+        data = json.load(mapfile)
+
+        self.latch_count = data["latch_count"]
+        self.input_count = data["input_count"]
+
+        self.clocks = data["clocks"]
+
+        self.sigmap = WitnessSigMap()
+        self.init_inputs = set(init["input"] for init in data["inits"])
+
+        for bit in data["inputs"] + data["seqs"] + data["inits"]:
+            self.sigmap.add_bit((tuple(bit["path"]), bit["offset"]), bit["input"])
+
+
+
+@cli.command(help="""
+Convert an AIGER witness trace into a Yosys witness trace.
+
+This requires a Yosys witness AIGER map file as generated by 'write_aiger -ywmap'.
+""")
+@click.argument("input", type=click.File("r"))
+@click.argument("mapfile", type=click.File("r"))
+@click.argument("output", type=click.File("w"))
+def aiw2yw(input, mapfile, output):
+    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}")
+    aiger_map = AigerMap(mapfile)
+
+    header_lines = list(itertools.islice(input, 0, 2))
+
+    if len(header_lines) == 2 and header_lines[1][0] in ".bcjf":
+        status = header_lines[0].strip()
+        if status == "0":
+            raise click.ClickException(f"{input_name}: file contains no trace, the AIGER status is unsat")
+        elif status == "2":
+            raise click.ClickException(f"{input_name}: file contains no trace, the AIGER status is sat")
+        elif status != "1":
+            raise click.ClickException(f"{input_name}: unexpected data in AIGER witness file")
+    else:
+        input = itertools.chain(header_lines, input)
+
+    ffline = next(input, None)
+    if ffline is None:
+        raise click.ClickException(f"{input_name}: unexpected end of file")
+    ffline = ffline.strip()
+    if not re.match(r'[01x]*$', ffline):
+        raise click.ClickException(f"{input_name}: unexpected data in AIGER witness file")
+    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')
+
+    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)
+
+    missing = set()
+
+    while True:
+        inline = next(input, None)
+        if inline is None:
+            click.echo(f"Warning: {input_name}: file may be incomplete")
+            break
+        inline = inline.strip()
+        if inline in [".", "# DONE"]:
+            break
+        if inline.startswith("#"):
+            continue
+
+        if not re.match(r'[01x]*$', ffline):
+            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")
+
+        values = WitnessValues()
+        for i, v in enumerate(inline):
+            if v == "x" or outyw.t > 0 and i in aiger_map.init_inputs:
+                continue
+
+            try:
+                bit = aiger_map.sigmap.id_to_bit[i]
+            except IndexError:
+                bit = None
+            if bit is None:
+                missing.insert(i)
+
+            values[bit] = v
+
+        outyw.step(values)
+
+    outyw.end_trace()
+
+    if missing:
+        click.echo("The following AIGER inputs belong to unknown signals:")
+        click.echo("  " + " ".join(str(id) for id in sorted(missing)))
+
+    click.echo(f"Converted {outyw.t} time steps.")
+
+@cli.command(help="""
+Convert a Yosys witness trace into an AIGER witness trace.
+
+This requires a Yosys witness AIGER map file as generated by 'write_aiger -ywmap'.
+""")
+@click.argument("input", type=click.File("r"))
+@click.argument("mapfile", type=click.File("r"))
+@click.argument("output", type=click.File("w"))
+def yw2aiw(input, mapfile, output):
+    click.echo(f"Converting Yosys witness trace {input.name!r} to AIGER witness trace {output.name!r}...")
+    click.echo(f"Using Yosys witness AIGER map file {mapfile.name!r}")
+    aiger_map = AigerMap(mapfile)
+    inyw = ReadWitness(input)
+
+    print("1", file=output)
+    print("b0", file=output)
+    # TODO the b0 status isn't really accurate, but we don't have any better info here
+    print("0" * aiger_map.latch_count, file=output)
+
+    all_missing = set()
+
+    for t, step in inyw.steps():
+        bits, missing = step.pack_present(aiger_map.sigmap)
+        bits = bits[::-1].replace('?', 'x')
+        all_missing.update(missing)
+        bits += 'x' * (aiger_map.input_count - len(bits))
+        print(bits, file=output)
+
+    print(".", file=output)
+
+    if all_missing:
+        click.echo("The following signals are missing in the AIGER map file and will be dropped:")
+        for sig in coalesce_signals(WitnessSig(*bit) for bit in all_missing):
+            click.echo("  " + sig.pretty())
+
+
+    click.echo(f"Converted {len(inyw)} time steps.")
+
 if __name__ == "__main__":
     cli()