symex: structural Concat-fold for byte-shadow load/store round-trips

The byte-shadow store/load pipeline used z3.simplify to recover load-
modify-store identity. z3 folds Extract(7, 0, X-Y) through byte
arithmetic (the low byte has no incoming borrow), but leaves the upper
Extracts unfolded because their borrow chain depends on the low byte.
After a round trip, the value was Concat(Extract(31, 8, X), 212 + 255*b)
— the parent identity for the low byte was destroyed and the full-Concat-
collapse rule could no longer recover X.

New _z3_concat_fold helper: flattens nested binary Concats z3 stores
internally into a single MSB→LSB sequence, merges adjacent same-parent
Extracts at contiguous bit ranges, collapses a full-width single Extract
to its parent. Pure structural rewrite — never folds Extract through
arithmetic.

_shadow_write and _shadow_read now use the structural fold instead of
z3.simplify. Multi-byte round trips recover the parent expression
cleanly.

Tests: new test_concat_fold.py covers the fold rules directly, the
round-trip via _shadow_write + _shadow_read, partial-overwrite preserving
minimal Concat, and an assertion that z3.simplify can't recover the
parent (so the structural fold is load-bearing).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: pgoodman

bindings/Python/symex/dispatch.py

@@ -305,13 +305,18 @@ def _shadow_write(shadow, addr, val, size):
     """
     z3 = _z3_module()
     if z3 is not None and _is_z3(val):
-        # Simplify before per-byte decomposition: when ``val`` was
-        # constructed by re-concatenating bytes earlier read from the
-        # shadow (the common round-trip), z3 collapses the
-        # ``Extract(i+7, i, Concat(b_n, …, b_0))`` patterns back to the
-        # original byte expressions. Without this, every load-modify-store
-        # cycle accumulates Concat/Extract layers around the same bytes.
-        val = z3.simplify(val)
+        # Structural Concat-fold before per-byte decomposition: when
+        # ``val`` was built by re-concatenating bytes earlier read from
+        # the shadow (the load-modify-store round-trip), the fold
+        # collapses ``Concat(Extract(31,24,X),Extract(23,16,X),...)``
+        # back to ``X`` (or a single Extract) so the per-byte
+        # decomposition below stores raw ``Extract(8i+7, 8i, X)``
+        # entries. Using the structural fold here instead of
+        # ``z3.simplify`` avoids the asymmetric-byte-arith collapse
+        # that would otherwise produce ``212 + 255*b`` for the low
+        # byte of ``1492 - zext_32(b)`` and lose the parent identity
+        # for future loads.
+        val = _z3_concat_fold(z3, val)
         i = 0
         while i < size:
             shadow[addr + i] = _z3_byte_at(val, i)
@@ -369,7 +374,13 @@ def _shadow_read(shadow, addr, size, data, buf):
     while j >= 0:
         result = z3.Concat(result, buf[j])
         j -= 1
-    return z3.simplify(result)
+    # Structural fold only — collapses consecutive Extracts of the
+    # same parent back to the parent (or a single Extract). Avoids
+    # ``z3.simplify`` here on purpose: simplify folds the low byte
+    # through byte arith asymmetrically and prevents the parent from
+    # being recovered. A consumer op (compare, branch, cast) is the
+    # right place to run a full simplify.
+    return _z3_concat_fold(z3, result)
 
 
 def _make_default_mem_read(is_float, shadow=None, buf=None, byte_order="little"):
@@ -569,6 +580,82 @@ def _z3_byte_at(val, i):
     return z3.Extract(8 * (i + 1) - 1, 8 * i, val)
 
 
+def _z3_concat_fold(z3, expr):
+    """Structural fold for Concat-of-consecutive-Extracts-of-same-parent.
+
+    Pure structure rewrite — never folds Extract through arithmetic.
+    `z3.simplify` would do that asymmetrically (the low byte of a
+    subtraction has no incoming borrow, so it folds into byte arith,
+    while high bytes don't), producing shapes like
+    `Concat(Extract(31, 8, X), 212 + 255*b)` where the low byte's
+    parent identity has been lost and the round-trip back to X is
+    impossible to recover.
+
+    Steps:
+
+      1. Flatten nested Concats — z3 stores n-ary Concat as a
+         left-associated chain of binary applications, so
+         `Concat(a, b, c, d)` is `Concat(Concat(Concat(a, b), c), d)`
+         under the hood. The fold has to see it as a flat MSB→LSB
+         sequence to merge consecutive Extracts.
+      2. Merge consecutive Extracts of the same parent at adjacent
+         bit ranges into a single Extract.
+      3. If a single Extract spans the parent's full width, return
+         the parent itself.
+
+    Rules 2 + 3 together recover the original parent expression on a
+    full-width round trip: store a 32-bit value as four byte shadows,
+    load all four, and the result is the original 32-bit term — no
+    z3.simplify required, and (more importantly) no asymmetric
+    arithmetic folding to the low byte.
+    """
+    if not z3.is_app(expr) or expr.decl().kind() != z3.Z3_OP_CONCAT:
+        return expr
+
+    # Step 1: flatten nested Concats into a single MSB→LSB list.
+    parts = []
+    stack = [expr]
+    while stack:
+        node = stack.pop()
+        if z3.is_app(node) and node.decl().kind() == z3.Z3_OP_CONCAT:
+            # Children are MSB→LSB; push in reverse so the pop order
+            # walks them MSB-first.
+            kids = node.children()
+            for k in reversed(kids):
+                stack.append(k)
+        else:
+            parts.append(node)
+
+    if len(parts) == 1:
+        return parts[0]
+
+    # Step 2: left-to-right merge of adjacent same-parent Extracts.
+    out = []
+    for p in parts:
+        if (out and
+                z3.is_app(p) and p.decl().kind() == z3.Z3_OP_EXTRACT and
+                z3.is_app(out[-1]) and
+                out[-1].decl().kind() == z3.Z3_OP_EXTRACT and
+                p.arg(0).eq(out[-1].arg(0))):
+            top_hi, top_lo = out[-1].params()
+            bot_hi, bot_lo = p.params()
+            if top_lo == bot_hi + 1:
+                out[-1] = z3.Extract(top_hi, bot_lo, p.arg(0))
+                continue
+        out.append(p)
+
+    if len(out) == 1:
+        only = out[0]
+        # Step 3: full-width Extract collapses to its parent.
+        if z3.is_app(only) and only.decl().kind() == z3.Z3_OP_EXTRACT:
+            hi, lo = only.params()
+            parent = only.arg(0)
+            if lo == 0 and hi == parent.size() - 1:
+                return parent
+        return only
+    return z3.Concat(*out)
+
+
 # Per-arity dispatch tables: each entry is `(opcode_range, builder)`.
 # Builders for ops that need a z3 module function take (z3, a[, b]) so
 # the cached `_z3_module()` reference flows through; the rest are pure

tests/symex/test_concat_fold.py

@@ -0,0 +1,122 @@
+# Copyright (c) 2026-present, Trail of Bits, Inc.
+#
+# This source code is licensed in accordance with the terms specified in
+# the LICENSE file found in the root directory of this source tree.
+
+"""Structural Concat-fold and shadow round-trip tests.
+
+The shadow store/load pipeline must recover the original parent
+expression on a full-width round trip. `z3.simplify` can't do this —
+it folds the low byte through byte arithmetic asymmetrically (the low
+byte of `1492 - zext_32(b)` becomes `212 + 255*b`), which destroys
+the parent identity for the high bytes. The structural Concat-fold
+in `dispatch._z3_concat_fold` keeps the parent reachable.
+"""
+
+import pytest
+
+z3 = pytest.importorskip("z3")
+
+from multiplier.symex.dispatch import (
+    _z3_concat_fold, _shadow_write, _shadow_read, _z3_byte_at,
+)
+
+
+def _bv(val, w):
+    return z3.BitVecVal(val, w)
+
+
+# ---------------------------------------------------------------------------
+# Direct fold tests: structural rewrites only.
+# ---------------------------------------------------------------------------
+
+def test_concat_fold_full_width_round_trip():
+    """Concat of all four byte-Extracts of a 32-bit term collapses to
+    the term itself (rule: full-coverage Extract → identity)."""
+    b = z3.BitVec("b", 8)
+    x = z3.BitVecVal(1492, 32) - z3.ZeroExt(24, b)
+    parts = [z3.Extract(8 * (i + 1) - 1, 8 * i, x) for i in range(4)]
+    # MSB-first concat (matching _shadow_read's layout).
+    expr = z3.Concat(*reversed(parts))
+    out = _z3_concat_fold(z3, expr)
+    assert out.eq(x), f"expected parent recovery, got {out}"
+
+
+def test_concat_fold_partial_width_to_extract():
+    """Concat of two consecutive Extracts of the same parent collapses
+    to one Extract spanning both."""
+    x = z3.BitVec("x", 32)
+    expr = z3.Concat(z3.Extract(15, 8, x), z3.Extract(7, 0, x))
+    out = _z3_concat_fold(z3, expr)
+    assert out.eq(z3.Extract(15, 0, x)), f"got {out}"
+
+
+def test_concat_fold_mixed_parents_stays_split():
+    """Different parents → fold leaves the Concat alone."""
+    x = z3.BitVec("x", 32)
+    y = z3.BitVec("y", 32)
+    expr = z3.Concat(z3.Extract(31, 8, y), z3.Extract(7, 0, x))
+    out = _z3_concat_fold(z3, expr)
+    # Must be a Concat with two operands (not collapsed).
+    assert z3.is_app(out) and out.decl().kind() == z3.Z3_OP_CONCAT
+    assert out.num_args() == 2
+
+
+def test_concat_fold_simplify_does_NOT_recover_parent():
+    """Demonstrates the bug `_z3_concat_fold` exists to avoid: with
+    `z3.simplify`, the low byte is folded through arithmetic and the
+    full-Concat-collapse rule no longer matches."""
+    b = z3.BitVec("b", 8)
+    x = z3.BitVecVal(1492, 32) - z3.ZeroExt(24, b)
+    parts = [z3.Extract(8 * (i + 1) - 1, 8 * i, x) for i in range(4)]
+    expr = z3.Concat(*reversed(parts))
+    # Naive z3.simplify: bytes 1-3 collapse, byte 0 turns into byte arith.
+    naive = z3.simplify(expr)
+    assert not naive.eq(x), \
+        "z3.simplify unexpectedly recovered the parent — the asymmetric" \
+        " byte-fold isn't firing in this z3 build, so this test no" \
+        " longer guards what it was written for."
+    # Our structural fold recovers the parent regardless.
+    structural = _z3_concat_fold(z3, expr)
+    assert structural.eq(x), f"structural fold lost the parent: {structural}"
+
+
+# ---------------------------------------------------------------------------
+# Shadow round-trip: STORE + LOAD via the actual helpers.
+# ---------------------------------------------------------------------------
+
+def test_shadow_round_trip_recovers_parent_32bit():
+    """Storing a 32-bit symbolic term then reading it back yields the
+    same term, not a byte-split residue."""
+    b = z3.BitVec("b", 8)
+    x = z3.BitVecVal(1492, 32) - z3.ZeroExt(24, b)
+
+    shadow = {}
+    _shadow_write(shadow, addr=0x100, val=x, size=4)
+    # Concrete fallback bytes — _shadow_read prefers shadow entries
+    # over these, so the actual values don't matter for this test.
+    data = bytes(4)
+    out = _shadow_read(shadow, addr=0x100, size=4, data=data, buf=[])
+    assert out is not None
+    assert out.eq(x), f"round trip lost parent identity: {out}"
+
+
+def test_shadow_partial_overwrite_keeps_minimal_concat():
+    """Overwriting byte 0 of a 4-byte symbolic term with a new
+    8-bit value yields `Concat(Extract(31, 8, X), new_byte)` — the
+    high bytes still reference X, the low byte is the new value, and
+    no further fold collapses across the parent boundary."""
+    b = z3.BitVec("b", 8)
+    new_byte = z3.BitVec("new_byte", 8)
+    x = z3.BitVecVal(1492, 32) - z3.ZeroExt(24, b)
+
+    shadow = {}
+    _shadow_write(shadow, addr=0x100, val=x, size=4)
+    # Now overwrite byte 0 with a fresh symbolic byte.
+    _shadow_write(shadow, addr=0x100, val=new_byte, size=1)
+
+    data = bytes(4)
+    out = _shadow_read(shadow, addr=0x100, size=4, data=data, buf=[])
+    assert out is not None
+    expected = z3.Concat(z3.Extract(31, 8, x), new_byte)
+    assert out.eq(expected), f"got {out}, expected {expected}"