Fix dmap cfgs.

2 files changed, 4 insertions, 1 deletions

diff --git a/pyecsca/sca/re/tree.py b/pyecsca/sca/re/tree.py
index 67637ee..f19efac 100644
--- a/pyecsca/sca/re/tree.py
+++ b/pyecsca/sca/re/tree.py
@@ -146,7 +146,7 @@ class Map:
 
     @property
     def cfgs(self) -> Set[Any]:
-        return set().union(*self.cfg_map.index)
+        return set(self.cfg_map.index)
 
     def __getitem__(self, item):
         if isinstance(item, tuple):
@@ -275,6 +275,8 @@ class Tree:
                 f"Dmaps: {len(self.maps)}",
                 f"Total cfgs: {len(self.root.cfgs)}",
                 f"Depth: {self.height}",
+                f"Size: {self.size}",
+                f"Leaves: {len(leaf_sizes)}",
                 f"Leaf sizes: {sorted(leaf_sizes)}",
                 f"Average leaf size: {np.mean(leaf_sizes):.3}",
                 f"Mean result size: {np.mean(leafs):.3}",
diff --git a/test/sca/test_tree.py b/test/sca/test_tree.py
index 234732d..a9638c5 100644
--- a/test/sca/test_tree.py
+++ b/test/sca/test_tree.py
@@ -13,6 +13,7 @@ def test_map():
     assert dmap.codomain == {True, False}
     assert dmap.mapping.index.tolist() == [0, 1]
     assert set(dmap.cfg_map.index) == cfgs
+    assert dmap.cfgs == cfgs
 
     io_map = {"a": {1: 5, 2: 7}, "b": {1: 3}}
     dmap = Map.from_io_maps(cfgs, io_map)