Fix HDF5 tracesets to allow more than 8k traces.

3 files changed, 36 insertions, 41 deletions

diff --git a/pyecsca/sca/trace_set/hdf5.py b/pyecsca/sca/trace_set/hdf5.py
index a3e72aa..caaba0b 100644
--- a/pyecsca/sca/trace_set/hdf5.py
+++ b/pyecsca/sca/trace_set/hdf5.py
@@ -73,7 +73,8 @@ class HDF5TraceSet(TraceSet):
         self._file = _file
         if _ordering is None:
             _ordering = [str(uuid.uuid4()) for _ in traces]
-        super().__init__(*traces, **kwargs, _ordering=_ordering)
+        self._ordering = _ordering
+        super().__init__(*traces, **kwargs)
 
     @classmethod
     def read(cls, input: Union[str, Path, bytes, BinaryIO]) -> "HDF5TraceSet":
@@ -83,17 +84,16 @@ class HDF5TraceSet(TraceSet):
             hdf5 = h5py.File(input, mode="r")
         else:
             raise TypeError
-        kwargs = dict(hdf5.attrs)
-        kwargs["_ordering"] = (
-            list(kwargs["_ordering"]) if "_ordering" in kwargs else list(hdf5.keys())
-        )
+        group = hdf5["traces"]
+        kwargs = dict(group.attrs)
+        ordering = []
         traces = []
-        for k in kwargs["_ordering"]:
-            meta = dict(HDF5Meta(hdf5[k].attrs))
-            samples = hdf5[k]
+        for k, samples in group.items():
+            meta = dict(HDF5Meta(samples.attrs))
             traces.append(Trace(np.array(samples, dtype=samples.dtype), meta))
+            ordering.append(k)
         hdf5.close()
-        return HDF5TraceSet(*traces, **kwargs)
+        return HDF5TraceSet(*traces, **kwargs, _ordering=ordering)
 
     @classmethod
     def inplace(cls, input: Union[str, Path, bytes, BinaryIO]) -> "HDF5TraceSet":
@@ -103,38 +103,34 @@ class HDF5TraceSet(TraceSet):
             hdf5 = h5py.File(input, mode="a")
         else:
             raise TypeError
-        kwargs = dict(hdf5.attrs)
-        kwargs["_ordering"] = (
-            list(kwargs["_ordering"]) if "_ordering" in kwargs else list(hdf5.keys())
-        )
+        group = hdf5["traces"]
+        kwargs = dict(group.attrs)
+        ordering = []
         traces = []
-        for k in kwargs["_ordering"]:
-            meta = HDF5Meta(hdf5[k].attrs)
-            samples = hdf5[k]
+        for k, samples in group.items():
+            meta = HDF5Meta(samples.attrs)
             traces.append(Trace(samples, meta))
-        return HDF5TraceSet(*traces, **kwargs, _file=hdf5)  # type: ignore[misc]
+            ordering.append(k)
+        return HDF5TraceSet(*traces, **kwargs, _file=hdf5, _ordering=ordering)  # type: ignore[misc]
 
-    def insert(self, index: int, value: Trace) -> Trace:
+    def append(self, value: Trace) -> Trace:
         key = str(uuid.uuid4())
-        self._ordering.insert(index, key)
         if self._file is not None:
-            new_samples = self._file.create_dataset(key, data=value.samples)
+            group = self._file["traces"]
+            new_samples = group.create_dataset(key, data=value.samples)
             new_meta = HDF5Meta(new_samples.attrs)
             if value.meta:
                 for k, v in value.meta.items():
                     new_meta[k] = v
             value = Trace(new_samples, new_meta)
-            self._file.attrs["_ordering"] = self._ordering
 
-        self._traces.insert(index, value)
+        self._ordering.append(key)
+        self._traces.append(value)
         return value
 
     def get(self, index: int) -> Trace:
         return self[index]
 
-    def append(self, value: Trace) -> Trace:
-        return self.insert(len(self), value)
-
     def remove(self, value: Trace):
         if value in self._traces:
             index = self._traces.index(value)
@@ -142,8 +138,8 @@ class HDF5TraceSet(TraceSet):
             self._ordering.remove(key)
             self._traces.remove(value)
             if self._file:
-                self._file.pop(key)
-                self._file.attrs["_ordering"] = self._ordering
+                group = self._file["traces"]
+                group.pop(key)
         else:
             raise KeyError
 
@@ -173,12 +169,11 @@ class HDF5TraceSet(TraceSet):
             hdf5 = h5py.File(output, "w")
         else:
             raise ValueError
+        group = hdf5.create_group("traces", track_order=True)
         for k in self._keys:
-            hdf5.attrs[k] = getattr(self, k)
-        hdf5.attrs["_ordering"] = self._ordering
-        for i, k in enumerate(self._ordering):
-            trace = self[i]
-            dset = hdf5.create_dataset(k, data=trace.samples)
+            group.attrs[k] = getattr(self, k)
+        for k, trace in zip(self._ordering, self):
+            dset = group.create_dataset(k, data=trace.samples)
             if trace.meta:
                 meta = HDF5Meta(dset.attrs)
                 for key, val in trace.meta.items():
@@ -201,4 +196,4 @@ class HDF5TraceSet(TraceSet):
                 if not key.startswith("_")
             ]
         )
-        return f"HDF5TraceSet('{fname}'{status}, {args})"
+        return f"HDF5TraceSet('{fname}'{status}, {args} <{len(self)}>)"
diff --git a/test/data/sca/test.h5 b/test/data/sca/test.h5
index f51b5fd..dc44b9d 100644
--- a/test/data/sca/test.h5
+++ b/test/data/sca/test.h5
diff --git a/test/sca/test_traceset.py b/test/sca/test_traceset.py
index a338a9d..ebc98a1 100644
--- a/test/sca/test_traceset.py
+++ b/test/sca/test_traceset.py
@@ -117,20 +117,20 @@ def test_h5_inplace():
         shutil.copy(orig_path, path)
         trace_set = HDF5TraceSet.inplace(path)
         assert trace_set is not None
-        test_trace = Trace(np.array([6, 7], dtype=np.dtype("i1")), meta={"thing": "ring"})
-        other_trace = Trace(np.array([15, 7], dtype=np.dtype("i1")), meta={"a": "b"})
+        test_trace = Trace(np.array([4, 7], dtype=np.dtype("i1")), meta={"thing": "ring"})
+        other_trace = Trace(np.array([8, 7], dtype=np.dtype("i1")), meta={"a": "b"})
         trace_set.append(test_trace)
-        assert len(trace_set) == 4
+        assert len(trace_set) == 3
         trace_set.append(other_trace)
         trace_set.remove(other_trace)
-        assert len(trace_set) == 4
+        assert len(trace_set) == 3
         trace_set.save()
         trace_set.close()
         test_set = HDF5TraceSet.read(path)
-        assert test_set.get(3) == test_set[3]
-        assert np.array_equal(test_set[3].samples, test_trace.samples)
-        assert test_set[3].meta["thing"] == test_trace.meta["thing"]
-        assert test_set[3] == test_trace
+        assert test_set.get(2) == test_set[2]
+        assert np.array_equal(test_set[2].samples, test_trace.samples)
+        assert test_set[2].meta["thing"] == test_trace.meta["thing"]
+        assert test_set[2] == test_trace
 
 
 def test_h5_save(example_traces, example_kwargs):