vcd improvements

src/kyupy/vcd.py

@@ -6,7 +6,7 @@ All variable values are flattened. Offsets are stored in each variable metadata.
 """
 
 from collections import namedtuple
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 from lark import Lark, Transformer
 import numpy as np
@@ -21,7 +21,7 @@ class Var:
     idcode: str
     reference: str
     scope: 'Scope'
-    offset: int = 0
+    locs: list[int] = field(default_factory=list)
 
 
 class Scope:
@@ -83,18 +83,18 @@ class VcdHeader:
 
 
 class VcdVarMap:
-    def __init__(self, header: VcdHeader, filter = lambda var: True) -> None:
+    def __init__(self, header: VcdHeader, var_locs = lambda _: []) -> None:
         self.var_list: list[Var] = []
         def collect(scope):
-            self.var_list.extend(v for v in scope.vars if filter(v))
+            for v in scope.vars:
+                v.locs = var_locs(v)
+                if v.locs is not None and len(v.locs) > 0:
+                    assert len(v.locs) == v.width, f'{v.reference}: len(locs)={len(v.locs)} != width={v.width}'
+                    self.var_list.append(v)
             for sub in scope.sub_scopes:
                 collect(sub)
         collect(header.root_scope)
-        offset = 0
-        for var in self.var_list:
-            var.offset = offset
-            offset += var.width
-        self.total_width = offset
+        self.total_width = max((max(v.locs) for v in self.var_list), default=-1) + 1
         self.idcode2var = {var.idcode: var for var in self.var_list}
 
 
@@ -149,12 +149,12 @@ GRAMMAR = r"""
     """
 
 
-def load(file, var_filter = lambda var: True, step_filter = lambda time, values, var_map: True):
+def load(file, var_locs = lambda _: [], step_filter = lambda *_: True):
     """Parses the contents of ``file`` as Verilog Change Dump (VCD).
 
     :param file: A file name or a file handle. Files with `.gz`-suffix are decompressed on-the-fly.
-    :param var_filter: A callback function to include only certain variables in the output.
-    :param step_filter: A callback function to include only certain steps in the output.
+    :param var_locs: A callback ``(var) -> list[int]`` mapping each variable to ndarray column indices. Empty list drops the variable.
+    :param step_filter: A callback ``(time, values, var_map) -> bool`` to select which timesteps to include.
     :return: A VcdData object with metadata and an ndarray with all values.
     """
     vcd = readtext(file)
@@ -163,7 +163,7 @@ def load(file, var_filter = lambda var: True, step_filter = lambda time, values,
     vcd_header_str = vcd[:header_size]
     vcd_header : VcdHeader = Lark(GRAMMAR, parser="lalr", lexer='contextual', transformer=VcdHeaderTransformer()).parse(vcd_header_str) # type: ignore
     vcd_data = vcd[header_size:].splitlines()
-    var_map = VcdVarMap(vcd_header, var_filter)
+    var_map = VcdVarMap(vcd_header, var_locs)
 
     chunk_size = 10240
     chunks = []
@@ -203,8 +203,9 @@ def load(file, var_filter = lambda var: True, step_filter = lambda time, values,
         if len(value_str) < var.width:
             value_str = _pad_char.get(value_str[0], '0') * (var.width - len(value_str)) + value_str
         for i, c in enumerate(value_str):
-            chunk[step_idx, var.offset + var.width - 1 - i] = _val_map.get(c, logic.UNKNOWN)
+            if var.locs[i] >= 0:
+                chunk[step_idx, var.locs[i]] = _val_map.get(c, logic.UNKNOWN)
 
     chunks.append(chunk[:step_idx])
-    data = np.concatenate(chunks, axis=0)
+    data = np.concatenate(chunks, axis=0).T
     return VcdData(var_map, steps, data)