Add more utilities for working with maps and trees.

8 files changed, 240 insertions, 102 deletions

diff --git a/pyecsca/ec/curve.py b/pyecsca/ec/curve.py
index 380ac2b..c952bd9 100644
--- a/pyecsca/ec/curve.py
+++ b/pyecsca/ec/curve.py
@@ -350,4 +350,4 @@ class EllipticCurve:
 
     def __repr__(self):
         params = ", ".join((f"{key}={val}" for key, val in self.parameters.items()))
-        return f"{self.__class__.__name__}([{params}] on {self.model} using {self.coordinate_model})"
+        return f"{self.__class__.__name__}([{params}] p={self.prime} on {self.coordinate_model})"
diff --git a/pyecsca/ec/mult/comb.py b/pyecsca/ec/mult/comb.py
index 252abd7..860dab7 100644
--- a/pyecsca/ec/mult/comb.py
+++ b/pyecsca/ec/mult/comb.py
@@ -60,6 +60,14 @@ class BGMWMultiplier(AccumulatorMultiplier, ScalarMultiplier):
         self.direction = direction
         self.width = width
 
+    def __hash__(self):
+        return id(self)
+
+    def __eq__(self, other):
+        if not isinstance(other, BGMWMultiplier):
+            return False
+        return self.formulas == other.formulas and self.short_circuit == other.short_circuit and self.width == other.width and self.direction == other.direction and self.accumulation_order == other.accumulation_order
+
     def __repr__(self):
         return f"{self.__class__.__name__}({', '.join(map(str, self.formulas.values()))}, short_circuit={self.short_circuit}, width={self.width}, direction={self.direction.name}, accumulation_order={self.accumulation_order.name})"
 
@@ -138,6 +146,14 @@ class CombMultiplier(AccumulatorMultiplier, ScalarMultiplier):
         )
         self.width = width
 
+    def __hash__(self):
+        return id(self)
+
+    def __eq__(self, other):
+        if not isinstance(other, CombMultiplier):
+            return False
+        return self.formulas == other.formulas and self.short_circuit == other.short_circuit and self.width == other.width and self.accumulation_order == other.accumulation_order
+
     def __repr__(self):
         return f"{self.__class__.__name__}({', '.join(map(str, self.formulas.values()))}, short_circuit={self.short_circuit}, width={self.width}, accumulation_order={self.accumulation_order.name})"
 
diff --git a/pyecsca/ec/params.py b/pyecsca/ec/params.py
index 3a6a54d..d1ee4f8 100644
--- a/pyecsca/ec/params.py
+++ b/pyecsca/ec/params.py
@@ -109,7 +109,7 @@ class DomainParameters:
         return f"{self.__class__.__name__}({name})"
 
     def __repr__(self):
-        return f"{self.__class__.__name__}({self.curve!r}, {self.generator!r}, {self.order}, {self.cofactor})"
+        return f"{self.__class__.__name__}({self.curve!r}, gen={self.generator!r}, ord={self.order}, cof={self.cofactor})"
 
 
 @public
diff --git a/pyecsca/ec/point.py b/pyecsca/ec/point.py
index bdbba5e..f7eb78d 100644
--- a/pyecsca/ec/point.py
+++ b/pyecsca/ec/point.py
@@ -211,7 +211,7 @@ class Point:
         return f"[{args}]"
 
     def __repr__(self):
-        return f"Point([{str(self)}] in {self.coordinate_model})"
+        return f"Point({str(self)} in {self.coordinate_model})"
 
 
 @public
diff --git a/pyecsca/sca/re/rpa.py b/pyecsca/sca/re/rpa.py
index 7bd796e..df36ea7 100644
--- a/pyecsca/sca/re/rpa.py
+++ b/pyecsca/sca/re/rpa.py
@@ -9,7 +9,7 @@ from typing import MutableMapping, Optional, Callable, List
 
 from sympy import FF, sympify, Poly, symbols
 
-from .tree import build_distinguishing_tree
+from .tree import Tree, Map
 from ...ec.coordinates import AffineCoordinateModel
 from ...ec.formula import (
     FormulaAction,
@@ -263,21 +263,18 @@ def rpa_distinguish(
                 used |= multiply_multiples
             mults_to_multiples[mult] = used
 
-        tree = build_distinguishing_tree(set(mults), mults_to_multiples)
+        dmap = Map.from_binary_sets(set(mults), mults_to_multiples)
+
+        tree = Tree.build(set(mults), dmap)
         log("Built distinguishing tree.")
-        log(
-            RenderTree(tree).by_attr(
-                lambda n: n.name
-                if n.name
-                else [mult.__class__.__name__ for mult in n.cfgs]
-            )
-        )
-        if tree is None or not tree.children:
+        log()
+        #tree.render()
+        if tree.size == 1:
             tries += 1
             continue
-        current_node = tree
+        current_node = tree.root
         while current_node.children:
-            _, best_distinguishing_multiple = current_node.name
+            best_distinguishing_multiple = current_node.dmap_input
             P0_inverse = rpa_input_point(best_distinguishing_multiple, P0, params)
             responses = []
             for _ in range(majority):
@@ -295,7 +292,7 @@ def rpa_distinguish(
                     best_distinguishing_multiple in mults_to_multiples[mult],
                 )
             response_map = {
-                child.oracle_response: child for child in current_node.children
+                child.response: child for child in current_node.children
             }
             current_node = response_map[response]
             mults = current_node.cfgs
diff --git a/pyecsca/sca/re/tree.py b/pyecsca/sca/re/tree.py
index 306893c..37b9ed9 100644
--- a/pyecsca/sca/re/tree.py
+++ b/pyecsca/sca/re/tree.py
@@ -1,96 +1,200 @@
+"""
+Tools for working with distinguishing maps and trees.
+
+         ________
+     ,o88~~88888888o,
+   ,~~?8P  88888     8,
+  d  d88 d88 d8_88     b
+ d  d888888888          b
+ 8,?88888888  d8.b o.   8
+ 8~88888888~ ~^8888  db 8
+ ?  888888          ,888P
+  ?  `8888b,_      d888P
+   `   8888888b   ,888'
+     ~-?8888888 _.P-~
+         ~~~~~~
+
+Here we grow the trees.
+ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
+ ⠀⠀⠀⠀⠀⠀⠀⠀⢰⣆⠀⢿⢢⡀⠸⣱⡀⢀⢤⡶⡄⠀⠀⠀⠀⠀⠀⠀⠀
+ ⠀⠀⠀⢰⣤⡴⣴⠢⣧⣿⢤⣈⠺⡧⣺⠷⡟⠓⠃⣇⢧⣶⣀⠀⠀⠀⠀⠀⠀
+ ⠀⢀⡀⢠⠽⣇⣹⠓⠽⡌⠓⠹⢶⢿⠀⢠⠁⡈⡶⠓⠋⠙⢻⡹⢓⡶⠀⠀⠀
+ ⠀⣬⣿⡎⠓⠒⠻⢄⡀⢳⡀⠀⢸⠈⢦⢸⠀⣸⢇⠀⢀⣠⡞⢉⣡⣴⣲⠄⠀
+ ⢠⣽⣈⣇⣴⣲⠖⠀⠉⠚⠳⣄⢸⠀⠈⣿⠊⡼⡦⢏⠉⠀⠉⢙⠮⢱⡆⠀⠀
+ ⠘⢎⡇⠘⢧⡀⣀⣤⣤⠄⠀⠈⢻⣇⢀⣽⠖⢣⣣⡤⠤⠒⠚⣝⣆⠀⠀⠀⠀
+ ⠀⠀⠓⠢⠤⠽⢧⣀⠀⠀⠀⠀⠀⣿⠏⢹⡴⠋⠙⠒⠶⢖⢤⡀⢀⣤⣶⠗⠀
+ ⡢⢄⠐⠮⠷⢆⠀⠈⠙⠛⠛⠻⢶⣿⠀⣾⢀⣀⣀⣤⠤⡼⡓⢋⣿⠫⢕⡆⠀
+ ⠈⣻⠧⠤⠤⠤⠿⠗⠒⠒⠛⠶⣤⣿⣼⠛⠉⠙⠒⣶⡖⠳⡗⠺⡿⣤⣤⢄⡀
+ ⠰⣹⠀⠀⠀⢟⡆⠀⠀⠀⠀⠀⠈⣿⡿⠀⠀⠀⠀⢧⠇⠀⠀⠀⠀⠙⠮⡯⠀
+ ⠀⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣿⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
+ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⣿⣇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
+ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢹⣿⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
+ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠘⣿⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
+ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣿⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
+ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣿⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
+ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣿⠇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
+ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣼⣿⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
+ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢠⣤⠾⠛⣿⠙⠛⠶⢦⠄⠀⠀⠀⠀⠀⠀⠀⠀⠀
+"""
 from copy import deepcopy
-from typing import Mapping, Any, Set, List
+from typing import Mapping, Any, Set, List, Tuple, Optional
 from collections import Counter
+
+import numpy as np
+import pandas as pd
 from public import public
-from anytree import Node
+from anytree import RenderTree, NodeMixin, AbstractStyle
 
 
 @public
-def build_distinguishing_tree(
-    cfgs: Set[Any], *cfg2resp_maps: Mapping[Any, Set[Any]], **kwargs
-) -> Node:
-    """
-    Build a distinguishing tree for given mappings of configs to True oracle responses.
+class Map:
+    """A distinguishing map."""
+    context: Any
+    mapping: pd.DataFrame
 
-    :param cfgs: The configurations to distinguish.
-    :param cfg2resp_maps: The mappings of configs to sets of inputs that give True oracle responses.
-    :param kwargs: Additional keyword arguments that will be passed to the node.
-    :return: A distinguishing tree.
-    """
-    n_cfgs = len(cfgs)
+    def __init__(self, context: Any, mapping: pd.DataFrame):
+        self.context = context
+        self.mapping = mapping
+
+    @classmethod
+    def from_binary_sets(cls, cfgs: Set[Any], mapping: Mapping[Any, Set[Any]]):
+        cfgs_l = list(cfgs)
+        inputs_l = list(set().union(*mapping.values()))
+        data = [[elem in mapping[cfg] for elem in inputs_l] for cfg in cfgs_l]
+        return Map(None, pd.DataFrame(data, index=cfgs_l, columns=inputs_l))
+
+
+@public
+class Node(NodeMixin):
+    """A node in a distinguishing tree."""
+    cfgs: Set[Any]
+    dmap_index: Optional[int]
+    dmap_input: Optional[Any]
+    response: Optional[Any]
+
+    def __init__(self, cfgs: Set[Any], dmap_index: Optional[int] = None, dmap_input: Optional[Any] = None, response: Optional[Any] = None, parent=None, children=None):
+        self.cfgs = cfgs
+        self.dmap_index = dmap_index
+        self.dmap_input = dmap_input
+        self.response = response
+
+        self.parent = parent
+        if children:
+            self.children = children
+
+
+@public
+class Tree:
+    """A distinguishing tree."""
+    maps: List[Map]
+    root: Node
+
+    def __init__(self, root: Node, *maps: Map):
+        self.maps = list(maps)
+        self.root = root
+
+    @property
+    def leaves(self) -> Tuple[Node]:
+        return self.root.leaves
+
+    @property
+    def height(self) -> int:
+        return self.root.height
+
+    @property
+    def size(self) -> int:
+        return self.root.size
+
+    def render(self) -> None:
+        style = AbstractStyle("\u2502  ", "\u251c\u2500\u2500", "\u2514\u2500\u2500")
+
+        def _str(n: Node):
+            pre = f">{n.response} " if n.response is not None else ""
+            if n.is_leaf:
+                return pre + "\n".join(str(cfg) for cfg in n.cfgs)
+            else:
+                return pre + f"{n.dmap_index}({n.dmap_input})"
+        print(RenderTree(self.root, style=style).by_attr(_str))
+
+    def describe(self) -> None:
+        leaf_sizes = [len(leaf.cfgs) for leaf in self.leaves]
+        print(f"Total cfgs: {len(self.root.cfgs)}")
+        print(f"Depth: {self.height}")
+        print(f"Leaf sizes: {sorted(leaf_sizes)}")
+        print(f"Average leaf size: {np.mean(leaf_sizes):.3}")
+        leafs = sum(([size] * size for size in leaf_sizes), [])
+        print(f"Mean result size: {np.mean(leafs):.3}")
+
+    def expand(self, dmap: Map) -> "Tree":
+        tree = deepcopy(self)
+        tree.maps.append(dmap)
+        for leaf in tree.leaves:
+            expanded = _build_tree(leaf.cfgs, *tree.maps, response=leaf.response)
+            # If we were able to split the leaf further, then replace it with the found tree.
+            if not expanded.is_leaf:
+                parent = leaf.parent
+                leaf.parent = None
+                expanded.parent = parent
+        return tree
+
+    @classmethod
+    def build(cls, cfgs: Set[Any], *maps: Map) -> "Tree":
+        return cls(_build_tree(cfgs, *maps), *maps)
+
+
+def _degree(split: pd.Series) -> float:
+    return len(split)
+
+
+def _size_of_largest(split: pd.Series) -> float:
+    return max(split)
+
+
+def _average_size(split: pd.Series) -> float:
+    return sum(split ** 2)/sum(split)
+
+
+def _build_tree(cfgs: Set[Any], *maps: Map, response: Optional[Any] = None) -> Node:
     # If there is only one remaining cfg, we do not need to continue and just return (base case 1).
-    # Note that n_cfgs will never be 0 here, as the base case 2 returns if the cfgs cannot be split into two sets (one would be empty).
+    # Note that n_cfgs will never be 0 here, as the base case 2 returns if the cfgs cannot
+    # be split into two sets (one would be empty).
+    n_cfgs = len(cfgs)
     if n_cfgs == 1:
-        return Node(None, cfgs=set(cfgs), **kwargs)
-
-    # Go over the maps and have a counter for each one
-    counters: List[Counter] = [Counter() for cfg2resp in cfg2resp_maps]
-    for counter, cfg2resp in zip(counters, cfg2resp_maps):
-        for cfg in cfgs:
-            elements = cfg2resp[cfg]
-            counter.update(elements)
+        return Node(set(cfgs), response=response)
 
-    nhalf = n_cfgs / 2
-    best_distinguishing_map = None
+    # Go over the maps and figure out which one splits the best.
     best_distinguishing_element = None
-    best_count = None
-    best_nhalf_distance = None
-    for i, (counter, cfg2resp) in enumerate(zip(counters, cfg2resp_maps)):
-        for multiple, count in counter.items():
-            if (
-                best_distinguishing_element is None
-                or abs(count - nhalf) < best_nhalf_distance
-            ):
-                best_distinguishing_map = i
-                best_distinguishing_element = multiple
-                best_count = count
-                best_nhalf_distance = abs(count - nhalf)
+    best_distinguishing_dmap = None
+    best_restricted = None
+    best_score = None
+    for dmap in maps:
+        # 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.filter(items=cfgs, axis=0)
+        for elem in dmap.mapping:
+            split = restricted[elem].value_counts()
+            # XXX: Try the other scores.
+            # TODO: Abort here when the total best score is already reached early.
+            score = _size_of_largest(split)
+            if best_score is None or score < best_score:
+                best_distinguishing_element = elem
+                best_distinguishing_dmap = dmap
+                best_score = score
+                best_restricted = restricted
 
+    # Now we have a dmap as well as an element in it that splits the best.
+    # Go over the groups of configs that share the response
+    groups = best_restricted.groupby(best_distinguishing_element).groups
     # We found nothing distinguishing the configs, so return them all (base case 2).
-    if best_count in (0, n_cfgs, None) or best_distinguishing_map is None:
-        return Node(None, cfgs=set(cfgs), **kwargs)
+    if len(groups) == 1:
+        return Node(set(cfgs), response=response)
 
-    result = Node(
-        (best_distinguishing_map, best_distinguishing_element), cfgs=set(cfgs), **kwargs
-    )
-    # Now go deeper and split based on the best-distinguishing element.
-    true_cfgs = set(
-        cfg
-        for cfg in cfgs
-        if best_distinguishing_element in cfg2resp_maps[best_distinguishing_map][cfg]
-    )
-    true_restricted_cfg2resps = [
-        {cfg: cfg2resp[cfg] for cfg in true_cfgs} for cfg2resp in cfg2resp_maps
-    ]
-    true_child = build_distinguishing_tree(
-        true_cfgs, *true_restricted_cfg2resps, oracle_response=True
-    )
-    true_child.parent = result
-
-    false_cfgs = cfgs.difference(true_cfgs)
-    false_restricted_cfg2resps = [
-        {cfg: cfg2resp[cfg] for cfg in false_cfgs} for cfg2resp in cfg2resp_maps
-    ]
-    false_child = build_distinguishing_tree(
-        false_cfgs, *false_restricted_cfg2resps, oracle_response=False
-    )
-    false_child.parent = result
-    return result
+    # Create our node
+    dmap_index = maps.index(best_distinguishing_dmap)
+    result = Node(set(cfgs), dmap_index, best_distinguishing_element, response=response)
 
+    for output, group_cfgs in groups.items():
+        child = _build_tree(set(group_cfgs), *maps, response=output)
+        child.parent = result
 
-def expand_tree(tree: Node, cfg2resp: Mapping[Any, Set[Any]]) -> Node:
-    """
-    Attempt to expand a given distinguishing tree with a new mapping of configs to True oracle responses.
-
-    :param tree: The tree to expand (will be copied).
-    :param cfg2resp: The new map.
-    :return: The expanded tree.
-    """
-    tree = deepcopy(tree)
-    for leaf in tree.leaves:
-        expanded = build_distinguishing_tree(leaf.cfgs, cfg2resp)
-        # If we were able to split the leaf further, then replace it with the found tree.
-        if not expanded.is_leaf:
-            leaf.name = expanded.name
-            leaf.children = expanded.children
-    return tree
+    return result
diff --git a/pyecsca/sca/re/zvp.py b/pyecsca/sca/re/zvp.py
index 2146966..9608745 100644
--- a/pyecsca/sca/re/zvp.py
+++ b/pyecsca/sca/re/zvp.py
@@ -202,6 +202,7 @@ def compute_factor_set(formula: Formula, affine: bool = True, filter_nonhomo: bo
 
     :param formula: Formula to compute the factor set of.
     :param affine: Whether to transform the polynomials into affine form.
+    :param filter_nonhomo: Whether to filter out non-homogenous polynomials.
     :return: The set of factors present in the formula.
     """
     unrolled = unroll_formula(formula)
@@ -566,19 +567,23 @@ def precomp_zvp_points(
     formulas: Mapping[str, Formula],
     params: DomainParameters,
     bound: int = 25,
+    filter_nonhomo: bool = True
 ) -> List[Mapping[Poly, Set[Point]]]:
     """
+    Precompute ZVP points for a given addition chain using given formula mapping.
 
-    :param chain:
-    :param formulas:
-    :param params:
-    :param bound:
-    :return:
+    :param chain: The addition chain.
+    :param formulas: Mapping of formula shortnames (e.g. "add") to formulas
+    :param params: The domain parameters to compute on.
+    :param bound: The bound for hard DCP cases.
+    :param filter_nonhomo: Whether to filter non-homogenous intermediate polynomials (will lead to false positives if False).
+    :return: A list with an entry for each formula application in the chain. Each entry is a mapping of intermediate polynomial to set of points.
     """
     factor_sets = {
-        formula: compute_factor_set(formula) for formula in formulas.values()
+        formula: compute_factor_set(formula, filter_nonhomo=filter_nonhomo) for formula in formulas.values()
     }
     # A bit of a hack to rename the poly variables for double as zvp_points expects that.
+    # TODO: Handle this better in the zvp_points function (to also be able to handle ladders and other stuff).
     for formula in formulas.values():
         if isinstance(formula, DoublingFormula):
             fset = factor_sets[formula]
diff --git a/test/sca/test_tree.py b/test/sca/test_tree.py
new file mode 100644
index 0000000..3dbb960
--- /dev/null
+++ b/test/sca/test_tree.py
@@ -0,0 +1,16 @@
+from pyecsca.sca.re.tree import Tree, Map
+import pandas as pd
+
+
+def test_build_tree():
+    cfgs = ["a", "b", "c"]
+    inputs1 = [1, 2, 3, 4]
+    mapping1 = pd.DataFrame([(0, 4, 5, 0), (0, 3, 0, 0), (1, 4, 0, 0)], columns=inputs1, index=cfgs)
+    dmap1 = Map(None, mapping1)
+
+    inputs2 = ["f", "e", "d"]
+    mapping2 = pd.DataFrame([(1, 0, 0), (2, 0, 0), (3, 0, 0)], columns=inputs2, index=cfgs)
+    dmap2 = Map(None, mapping2)
+    tree = Tree.build(set(cfgs), dmap1, dmap2)
+    print()
+    tree.render()
\ No newline at end of file