From fec53d95fc4d7deff67379029a36d1f7b82e1517 Mon Sep 17 00:00:00 2001
From: J08nY
Date: Fri, 12 Apr 2024 14:17:11 +0200
Subject: Fix tree building with deduplicated dmaps.

---
 pyecsca/sca/re/tree.py |  2 +-
 pyproject.toml         |  2 +-
 test/sca/test_tree.py  | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/pyecsca/sca/re/tree.py b/pyecsca/sca/re/tree.py
index 8a05987..f5d5b61 100644
--- a/pyecsca/sca/re/tree.py
+++ b/pyecsca/sca/re/tree.py
@@ -364,7 +364,7 @@ def _build_tree(
     for i, dmap in maps.items():
         # Now we have a map, it may be binary or have larger output domain
         # Note we should look at the restriction of the map to the current "cfgs" and split those
-        restricted = dmap.mapping.loc[dmap.cfg_map.loc[list(cfgs), "vals"].unique()]
+        restricted = dmap.mapping.loc[dmap.cfg_map.loc[list(cfgs), "vals"]]
         for j, column in restricted.items():
             split = column.value_counts(dropna=False)
             # TODO: Try the other scores.
diff --git a/pyproject.toml b/pyproject.toml
index 894fcd8..22c37d5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -65,7 +65,7 @@
 "leia" = ["smartleia"]
 "gmp" = ["gmpy2"]
 "pari" = ["cysignals", "cypari2"]
-"dev" = ["mypy", "flake8", "interrogate", "pyinstrument", "black", "types-setuptools"]
+"dev" = ["mypy", "flake8", "interrogate", "pyinstrument", "black", "types-setuptools", "pydocstyle"]
 "test" = ["pytest>=7.0.0", "coverage", "pytest-cov", "pytest-sugar", "pytest-mock"]
 "doc" = ["sphinx", "sphinx-autodoc-typehints", "nbsphinx", "sphinx-paramlinks", "sphinx_design", "alabaster>=0.7.16"]
 
diff --git a/test/sca/test_tree.py b/test/sca/test_tree.py
index 897ffd5..5db083c 100644
--- a/test/sca/test_tree.py
+++ b/test/sca/test_tree.py
@@ -1,4 +1,5 @@
 import random
+from collections import OrderedDict
 from copy import deepcopy
 
 from pyecsca.sca.re.tree import Tree, Map
@@ -90,6 +91,51 @@ def test_build_tree():
     tree.describe()
 
 
+def test_build_tree_dedup():
+    """Make sure that dmap deduplication does not affect the tree."""
+    cfgs = {"a", "b", "c", "d", "e", "f", "g"}
+    binary_sets = {
+        "a": {1, 2, 3},
+        "b": {2, 3, 4},
+        "c": {1, 2, 3},
+        "d": {4, 2},
+        "e": {4, 2},
+        "f": {4, 2},
+        "g": {4, 2},
+    }
+    dmap = Map.from_sets(cfgs, binary_sets)
+    original = deepcopy(dmap)
+    dmap.deduplicate()
+
+    tree = Tree.build(cfgs, original)
+    dedup = Tree.build(cfgs, dmap)
+    assert tree.describe() == dedup.describe()
+
+
+def test_build_tree_reorder():
+    """Make sure that dmap input order does not affect the tree."""
+    cfgs = {"a", "b", "c", "d", "e", "f", "g"}
+    binary_sets = {
+        "a": {1, 2, 3},
+        "b": {2, 3, 4},
+        "c": {1, 2, 3},
+        "d": {4, 2},
+        "e": {4, 2},
+        "f": {4, 2},
+        "g": {4, 2},
+    }
+    trees = set()
+    for i in range(10):
+        shuffled = list(binary_sets.items())
+        random.shuffle(shuffled)
+        bset = OrderedDict(shuffled)
+        dmap = Map.from_sets(cfgs, bset)
+        if i % 2 == 0:
+            dmap.deduplicate()
+        trees.add(Tree.build(cfgs, dmap).describe())
+    assert len(trees) == 1
+
+
 def test_expand_tree():
     cfgs = ["a", "b", "c"]
     cfg_map = pd.DataFrame([0, 1, 2], index=cfgs, columns=["vals"])
-- 
cgit v1.2.3-70-g09d2