From 18b44a1e84da3768056e23dca66ea62f36f044c3 Mon Sep 17 00:00:00 2001
From: Krystine Sherwin <93062060+KrystalDelusion@users.noreply.github.com>
Date: Mon, 22 May 2023 11:44:19 +1200
Subject: [PATCH] yosys-witness: add append option to yw2yw

Can now append a user defined number of steps to input traces when joining.
If the number of steps is +ve, inputs are all set to 0.
If -ve then steps are skipped.
If all of steps are skipped (including init step) then the input trace will not be copied.
If more than one input trace is provided, the append option will need to be provided the same number of times as there are input traces.
---
 backends/smt2/witness.py | 23 ++++++++++++++++++++---
 backends/smt2/ywio.py    | 20 ++++++++++++++++----
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/backends/smt2/witness.py b/backends/smt2/witness.py
index 8e13cba27..baac3176d 100644
--- a/backends/smt2/witness.py
+++ b/backends/smt2/witness.py
@@ -88,14 +88,31 @@ If two or more inputs are provided they will be concatenated together into the o
 """)
 @click.argument("inputs", type=click.File("r"), nargs=-1)
 @click.argument("output", type=click.File("w"))
-def yw2yw(inputs, output):
+@click.option("--append", "-p", type=int, multiple=True,
+              help="Number of steps (+ve or -ve) to append to end of input trace. "
+                  +"Can be defined multiple times, following the same order as input traces. ")
+def yw2yw(inputs, output, append):
     outyw = WriteWitness(output, "yosys-witness yw2yw")
     join_inputs = len(inputs) > 1
     inyws = {}
-    for input in inputs:
+
+    if not append:
+        # default to 0
+        append = [0] * len(inputs)
+    if len(append) != len(inputs):
+        print(f"Mismatch in number of --append values ({len(append)}) and input traces ({len(inputs)}).")
+        sys.exit(1)
+
+    for (input, p) in zip(inputs, append):
         if (join_inputs):
             click.echo(f"Loading signals from yosys witness trace {input.name!r}...")
         inyw = ReadWitness(input)
+        if p:
+            click.echo(f"  appending {p} steps")
+            if (p + len(inyw) <= 0):
+                click.echo(f"  skipping {input.name!r} (only {len(inyw)} steps to skip)")
+                continue
+            inyw.append_steps(p)
         inyws[input] = inyw
         for clock in inyw.clocks:
             if clock not in outyw.clocks:
@@ -119,7 +136,7 @@ def yw2yw(inputs, output):
         for t, values in inyw.steps(1):
             outyw.step(values)
 
-        click.echo(f"Copied {t + 1} time steps.")
+        click.echo(f"  copied {t + 1} time steps.")
         first_witness = False
 
     outyw.end_trace()
diff --git a/backends/smt2/ywio.py b/backends/smt2/ywio.py
index 2b897200f..4e95f8c33 100644
--- a/backends/smt2/ywio.py
+++ b/backends/smt2/ywio.py
@@ -393,12 +393,16 @@ class ReadWitness:
     def init_step(self):
         return self.step(0)
     
+    def non_init_bits(self):
+        if len(self) > 1:
+            return len(self.bits[1])
+        else:
+            return sum([sig.width for sig in self.signals if not sig.init_only])
+    
     def first_step(self):
         values = WitnessValues()
-        if len(self.bits) <= 1:
-            raise NotImplementedError("ReadWitness.first_step() not supported for less than 2 steps")
-        non_init_bits = len(self.bits[1])
-        values.unpack(WitnessSigMap([sig for sig in self.signals if not sig.init_only]), self.bits[0][-non_init_bits:])
+        # may have issues when non_init_bits is 0
+        values.unpack(WitnessSigMap([sig for sig in self.signals if not sig.init_only]), self.bits[0][-self.non_init_bits():])
         return values
 
     def step(self, t):
@@ -410,5 +414,13 @@ class ReadWitness:
         for i in range(start, len(self.bits)):
             yield i, self.step(i)
 
+    def append_steps(self, t):
+        if not t:
+            pass
+        elif t < 0:
+            self.bits = self.bits[:t]
+        else:
+            self.bits.extend(["0"*self.non_init_bits()]*t)
+
     def __len__(self):
         return len(self.bits)