Repository: simple-crypto/SCALib (language: python)

Issue:
Bug: assertion failed in bp_compute.rs
**Describe the bug**
In a graph where a VAR SINGLE is connected to both ends of a table and there are multiple executions, the program crashes with a failed assertion. The dimensions of a factor does not equal the dimension of the contained matrix. If an equivalent graph is created manually by replacing VAR MULTIPLE by VAR SINGLES, the error does not occur.

**To Reproduce**
This is a contrived example, but the same problem persists for larger graphs with more nodes between the input and output of the table. The package version used is v0.5.3.
```py
from scalib.attacks import FactorGraph, BPState
import numpy as np

graph_desc = f"""
NC 2
TABLE t
VAR SINGLE s

VAR MULTI a
VAR MULTI b

PROPERTY a = t[s]
PROPERTY b = a ^ s
"""

graph = FactorGraph(graph_desc, {"t": np.arange(2, dtype=np.uint32)})

sasca = BPState(graph, 2)

sasca.bp_loopy(4,initialize_states=True)
```

**Observed behavior**
```
thread '<unnamed>' panicked at 'assertion failed: `(left == right)`
  left: `(1, 2)`,
 right: `(2, 2)`', scalib/src/sasca/bp_compute.rs:158:17
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
Traceback (most recent call last):
  File "/tmp/test.py", line 21, in <module>
    sasca.bp_loopy(4, initialize_states=True)
  File "/home/stian/.local/lib/python3.10/site-packages/scalib/attacks/factor_graph.py", line 226, in bp_loopy
    self._inner.propagate_loopy_step(it)
pyo3_runtime.PanicException: assertion failed: `(left == right)`
  left: `(1, 2)`,
 right: `(2, 2)`
```
I.e. it panics at the assertion at https://github.com/simple-crypto/SCALib/blob/v0.5.3/src/scalib_ext/scalib/src/sasca/bp_compute.rs#L158

**Expected behavior**
The program should not crash and should behave the same as the manually created multiple execution graph:
```
NC 2
TABLE t
VAR SINGLE s

VAR SINGLE a0
VAR SINGLE b0

PROPERTY a0 = t[s]
PROPERTY b0 = a0 ^ s

VAR SINGLE a1
VAR SINGLE b1

PROPERTY a1 = t[s]
PROPERTY b1 = a1 ^ s
```
Which does not crash.



Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index c56a6459..d2f47bba 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -5,6 +5,8 @@ Changelog
 Not released
 ------------
 
+* Fix ``Distribution`` size bug in belief propagation computation (#102).
+
 v0.5.4 (2023/04/26)
 -------------------
 
diff --git a/src/scalib_ext/scalib/src/sasca/bp_compute.rs b/src/scalib_ext/scalib/src/sasca/bp_compute.rs
index b164bbfe..06c80d3e 100644
--- a/src/scalib_ext/scalib/src/sasca/bp_compute.rs
+++ b/src/scalib_ext/scalib/src/sasca/bp_compute.rs
@@ -227,12 +227,9 @@ impl Distribution {
             (DistrRepr::Full(vst), DistrRepr::Full(vdiv)) => (vst, vdiv),
         };
         res.value = DistrRepr::Full(
-            Zip::from(
-                vst.broadcast((std::cmp::max(vst.dim().0, vdiv.dim().0), vst.dim().1))
-                    .unwrap(),
-            )
-            .and_broadcast(vdiv)
-            .map_collect(|vst, vdiv| *vst / (*vdiv + MIN_PROBA)),
+            Zip::from(vst.broadcast((res.shape.0, vst.dim().1)).unwrap())
+                .and_broadcast(vdiv)
+                .map_collect(|vst, vdiv| *vst / (*vdiv + MIN_PROBA)),
         );
         return res;
     }


Test patch:
diff --git a/tests/test_factorgraph.py b/tests/test_factorgraph.py
index 75043100..09299532 100644
--- a/tests/test_factorgraph.py
+++ b/tests/test_factorgraph.py
@@ -936,3 +936,59 @@ def test_manytraces():
     distri_y_ref = distri_y_ref / distri_y_ref.sum(axis=1, keepdims=True)
 
     assert np.allclose(distri_y_ref, distri_y, rtol=1e-5, atol=1e-19)
+
+
+def test_mix_single_multi():
+    graph_desc = """
+    NC 2
+    TABLE t
+    VAR SINGLE s
+
+    VAR MULTI a
+    VAR MULTI b
+
+    PROPERTY TABLE: a = t[s]
+    PROPERTY XOR: b = a ^ s
+    """
+
+    graph_desc2 = """
+    NC 2
+    TABLE t
+    VAR SINGLE s
+
+    VAR SINGLE a0
+    VAR SINGLE b0
+
+    PROPERTY a0 = t[s]
+    PROPERTY b0 = a0 ^ s
+
+    VAR SINGLE a1
+    VAR SINGLE b1
+
+    PROPERTY a1 = t[s]
+    PROPERTY b1 = a1 ^ s
+    """
+
+    graph = FactorGraph(graph_desc, {"t": np.arange(2, dtype=np.uint32)})
+    graph2 = FactorGraph(graph_desc2, {"t": np.arange(2, dtype=np.uint32)})
+
+    sasca = BPState(graph, 2)
+    sasca.bp_loopy(4, initialize_states=True)  # check no crash
+
+    sasca = BPState(graph, 2)
+    sasca2 = BPState(graph2, 1)
+
+    for _ in range(10):
+        sasca.bp_loopy(1, initialize_states=False)
+        sasca2.bp_loopy(1, initialize_states=False)
+        for v in ("a", "b"):
+            d = sasca.get_distribution(v)
+            d0 = sasca2.get_distribution(f"{v}0")
+            d1 = sasca2.get_distribution(f"{v}0")
+            if d is not None:
+                assert np.allclose(d[0, :], d0, rtol=1e-5, atol=1e-19)
+                assert np.allclose(d[1, :], d1, rtol=1e-5, atol=1e-19)
+        d = sasca.get_distribution("s")
+        d2 = sasca2.get_distribution("s")
+        if d is not None:
+            assert np.allclose(d, d2, rtol=1e-5, atol=1e-19)