Add dmap merging.

2 files changed, 41 insertions, 1 deletions

diff --git a/pyecsca/sca/re/tree.py b/pyecsca/sca/re/tree.py
index 2893411..6aae1ea 100644
--- a/pyecsca/sca/re/tree.py
+++ b/pyecsca/sca/re/tree.py
@@ -141,6 +141,15 @@ class Map:
     def cfgs(self) -> Set[Any]:
         return set().union(*self.cfg_map.index)
 
+    def __getitem__(self, item):
+        if isinstance(item, tuple):
+            cfg, inp = item
+            row = self.cfg_map.loc[cfg, "vals"]
+            col = self.domain.index(inp)
+            return self.mapping.loc[row, col]
+        else:
+            raise KeyError
+
     def deduplicate(self):
         """Deduplicate the configs of this distinguishing map based on the rows."""
         for row, data in self.mapping.groupby(
@@ -150,7 +159,20 @@ class Map:
 
     def merge(self, other: "Map"):
         """Merge in another distinguishing map operating on different configs."""
-        pass
+        # Domains should be equal (but only as sets, we can reorder)
+        if set(self.domain) != set(other.domain):
+            raise ValueError("Cannot merge dmaps with different domains.")
+        reordering = [other.domain.index(elem) for elem in self.domain]
+        # Get the last used index in cfg_map
+        last = max(self.cfg_map["vals"])
+        # Offset the other cfg_map and mapping index by last + 1
+        other_cfg_map = other.cfg_map + (last + 1)
+        other_mapping = other.mapping[reordering].set_index(other.mapping.index + (last + 1))
+        # Now concat the cfg_map and mapping
+        self.cfg_map = pd.concat([self.cfg_map, other_cfg_map], copy=False)
+        self.mapping = pd.concat([self.mapping, other_mapping], copy=False)
+        # Finally, adjust the codomain
+        self.codomain.update(other.codomain)
 
 
 @public
diff --git a/test/sca/test_tree.py b/test/sca/test_tree.py
index 6ef83a8..ce16f97 100644
--- a/test/sca/test_tree.py
+++ b/test/sca/test_tree.py
@@ -20,6 +20,24 @@ def test_map():
     assert dmap.codomain == {5, 3, 7, None}
 
 
+def test_map_merge():
+    cfgs = {"a", "b"}
+    binary_sets = {"a": {1, 2, 3}, "b": {2, 4}}
+    dmap1 = Map.from_sets(cfgs, binary_sets)
+    assert len(dmap1.mapping) == 2
+
+    cfgs = {"c", "d"}
+    binary_sets = {"c": {1, 2}, "d": {2, 4, 3}}
+    dmap2 = Map.from_sets(cfgs, binary_sets)
+    assert len(dmap2.mapping) == 2
+    dmap1.merge(dmap2)
+    assert len(dmap1.mapping) == 4
+    assert len(dmap1.cfg_map) == 4
+    assert len(dmap1.codomain) == 2
+    assert not dmap1["c", 3]
+    assert dmap1["a", 0]
+
+
 def test_build_tree():
     cfgs = ["a", "b", "c"]
     cfg_map = pd.DataFrame([0, 1, 2], index=cfgs, columns=["vals"])