Merge pull request #3902 from YosysHQ/krys/yw_join

yosys-witness concat yw trace files

backends/smt2/witness.py

@@ -84,26 +84,65 @@ def stats(input):
 Transform a Yosys witness trace.
 
 Currently no transformations are implemented, so it is only useful for testing.
+If two or more inputs are provided they will be concatenated together into the output.
 """)
-@click.argument("input", type=click.File("r"))
+@click.argument("inputs", type=click.File("r"), nargs=-1)
 @click.argument("output", type=click.File("w"))
-def yw2yw(input, output):
-    click.echo(f"Copying yosys witness trace from {input.name!r} to {output.name!r}...")
-    inyw = ReadWitness(input)
+@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 clock in inyw.clocks:
-        outyw.add_clock(clock["path"], clock["offset"], clock["edge"])
+    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 sig in inyw.signals:
-        outyw.add_sig(sig.path, sig.offset, sig.width, sig.init_only)
+    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:
+                outyw.add_clock(clock["path"], clock["offset"], clock["edge"])
 
-    for t, values in inyw.steps():
-        outyw.step(values)
+        for sig in inyw.signals:
+            if sig not in outyw.signals:
+                outyw.add_sig(sig.path, sig.offset, sig.width, sig.init_only)
+
+    init_values = sum([inyw.init_step() for inyw in inyws.values()], start=WitnessValues())
+
+    first_witness = True
+    for (input, inyw) in inyws.items():
+        click.echo(f"Copying yosys witness trace from {input.name!r} to {output.name!r}...")
+
+        if first_witness:
+            outyw.step(init_values)
+        else:
+            outyw.step(inyw.first_step())
+
+        for t, values in inyw.steps(1):
+            outyw.step(values)
+
+        click.echo(f"  copied {t + 1} time steps.")
+        first_witness = False
 
     outyw.end_trace()
 
-    click.echo(f"Copied {outyw.t + 1} time steps.")
+    if join_inputs:
+        click.echo(f"Copied {outyw.t} total time steps.")
 
 
 class AigerMap:

backends/smt2/ywio.py

@@ -165,8 +165,8 @@ class WitnessSig:
         else:
             return f"{pretty_path(self.path)}[{self.offset}]"
 
-    def __eq__(self):
-        return self.sort_key
+    def __eq__(self, other):
+        return self.sort_key == other.sort_key
 
     def __hash__(self):
         return hash(self.sort_key)
@@ -294,6 +294,16 @@ class WitnessValues:
 
         return sorted(signals), missing_signals
 
+    def __add__(self, other: "WitnessValues"):
+        new = WitnessValues()
+        new += self
+        new += other
+        return new
+
+    def __iadd__(self, other: "WitnessValues"):
+        for key, value in other.values.items():
+            self.values.setdefault(key, value)
+        return self
 
 class WriteWitness:
     def __init__(self, f, generator):
@@ -380,14 +390,37 @@ class ReadWitness:
 
         self.bits = [step["bits"] for step in data["steps"]]
 
+    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()
+        # 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):
         values = WitnessValues()
         values.unpack(self.sigmap, self.bits[t])
         return values
 
-    def steps(self):
-        for i in range(len(self.bits)):
+    def steps(self, start=0):
+        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)