From 36f5bcf7ac87f14938254a1d67827ceda67c4192 Mon Sep 17 00:00:00 2001
From: hongyu <linerwang1997@gmail.com>
Date: Mon, 29 Jun 2026 15:28:29 +0100
Subject: [PATCH 01/11] FIX(decomp): grouping constraint leak; ordering

---
 .../region_decomp/__init__.py                 |  21 +-
 .../tests/test_grouped_region_decomp.py       | 407 +++++++++---------
 2 files changed, 209 insertions(+), 219 deletions(-)

diff --git a/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py b/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
index f25edc2c..7e41a99d 100644
--- a/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
+++ b/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
@@ -386,8 +386,10 @@ def _loop_group_regions(
     Invariants (across `reduce`/`loop` iterations):
     - `acc['regions']` shrinks monotonically: each iteration moves regions into
       `has` (grouped) or keeps them in `without` (remaining).
-    - `acc['constraint_path']` grows by one element per iteration (the current
-      `konstraint`), regardless of whether any regions matched.
+    - `acc['idx_path']` and `acc['constraint_path']` stay constant across
+      iterations (they describe this level). The current `konstraint` is only
+      prepended for the `has` branch's recursion, never carried into later
+      iterations, which process `without` regions that lack `konstraint`.
     - `acc['groups']` only grows: new groups are prepended when `has` is non-empty.
     """
 
@@ -424,8 +426,10 @@ def update_counter(s: str) -> None:
         return counter
 
     counter = mk_counter(all_constraints_with_dup)
-    assoc_list: list[tuple[str, int]] = [(k, v) for (k, v) in counter.items()]
-    sorted(assoc_list, key=lambda kv: kv[1], reverse=True)
+    # Most frequent first, ties broken alphabetically.
+    assoc_list: list[tuple[str, int]] = sorted(
+        counter.items(), key=lambda kv: (-kv[1], kv[0])
+    )
     constraints_by_most_frequent: list[str] = [kv[0] for kv in assoc_list]
 
     # grouped: tuple[list[RegionGroup], list[RegionStr]]
@@ -489,11 +493,16 @@ def loop(
             res = [group, *groups], without
         else:
             res = groups, without
+        # `idx_path` / `constraint_path` describe this level and must stay
+        # constant across reduce iterations. Only the `has` branch (recursed
+        # above) gets the extended `new_idx_path` / `new_constraint_path`; the
+        # `without` regions handled by later iterations do NOT contain
+        # `konstraint`, so it must not leak into their path.
         return Acc(
             groups=res[0],
             regions=res[1],
-            idx_path=new_idx_path,
-            constraint_path=new_constraint_path,
+            idx_path=idx_path,
+            constraint_path=constraint_path,
         )
 
     init = Acc(
diff --git a/packages/imandrax-api-models/tests/test_grouped_region_decomp.py b/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
index 0e057cc3..7c8a29cd 100644
--- a/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
+++ b/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
@@ -1,3 +1,5 @@
+import os
+
 import imandrax_api
 from imandrax_api.lib import RegionStr
 from inline_snapshot import snapshot
@@ -6,7 +8,6 @@
 
 
 def trust():
-    import os
     from typing import NoReturn
 
     import dotenv
@@ -31,7 +32,7 @@ def trust():
         if y > 0 then 5
         else 6"""
     _eval_res = c.eval_src(IML_CODE)
-    decomp_res = c.decompose(name='classify')
+    decomp_res = c.decompose(name='classify', string_results=True, prune=True)
 
     def raise_(exc: BaseException) -> NoReturn:
         raise exc
@@ -45,57 +46,58 @@ def raise_(exc: BaseException) -> NoReturn:
 
 
 def test():
-    # regions = trust()
-    # assert regions == snapshot()
-    regions = [
-        RegionStr(
-            constraints_str=['x <= 0', 'y <= 0'],
-            invariant_str='6',
-            model_str={'x': '0', 'y': '0'},
-            model_eval_str='6',
-        ),
-        RegionStr(
-            constraints_str=['y >= 1', 'x <= 0'],
-            invariant_str='5',
-            model_str={'x': '0', 'y': '1'},
-            model_eval_str='5',
-        ),
-        RegionStr(
-            constraints_str=['x >= 1', 'y >= (-10)', 'y <= 0'],
-            invariant_str='4',
-            model_str={'x': '1', 'y': '0'},
-            model_eval_str='4',
-        ),
-        RegionStr(
-            constraints_str=['x >= 1', 'y <= (-11)'],
-            invariant_str='3',
-            model_str={'x': '1', 'y': '(-11)'},
-            model_eval_str='3',
-        ),
-        RegionStr(
-            constraints_str=['x <= y', 'x >= 1', 'y >= 1'],
-            invariant_str='2',
-            model_str={'x': '1', 'y': '1'},
-            model_eval_str='2',
-        ),
-        RegionStr(
-            constraints_str=['not (x <= y)', 'x >= 1', 'y >= 1'],
-            invariant_str='1',
-            model_str={'x': '2', 'y': '1'},
-            model_eval_str='1',
-        ),
-    ]
+    regions = trust()
+    assert regions == snapshot(
+        [
+            RegionStr(
+                constraints_str=['y <= 0', 'x <= 0'],
+                invariant_str='6',
+                model_str={'x': '0', 'y': '0'},
+                model_eval_str='6',
+            ),
+            RegionStr(
+                constraints_str=['y >= 1', 'x <= 0'],
+                invariant_str='5',
+                model_str={'x': '0', 'y': '1'},
+                model_eval_str='5',
+            ),
+            RegionStr(
+                constraints_str=['y >= (-10)', 'y <= 0', 'x >= 1'],
+                invariant_str='4',
+                model_str={'x': '1', 'y': '0'},
+                model_eval_str='4',
+            ),
+            RegionStr(
+                constraints_str=['y <= (-11)', 'x >= 1'],
+                invariant_str='3',
+                model_str={'x': '1', 'y': '(-11)'},
+                model_eval_str='3',
+            ),
+            RegionStr(
+                constraints_str=['x <= y', 'y >= 1', 'x >= 1'],
+                invariant_str='2',
+                model_str={'x': '1', 'y': '1'},
+                model_eval_str='2',
+            ),
+            RegionStr(
+                constraints_str=['x > y', 'y >= 1', 'x >= 1'],
+                invariant_str='1',
+                model_str={'x': '2', 'y': '1'},
+                model_eval_str='1',
+            ),
+        ]
+    )
 
     hdr = HumDecomposeRes.from_regions(regions)
     assert hdr.to_tree_str() == snapshot("""\
-├── [1] constraint='x <= 0' invariant=None (w=2, n_children=2, n_descendants=2)
-│   ├── [1.1] constraint='y <= 0' invariant='6' (w=1, n_children=0, n_descendants=0)
-│   └── [1.1.2] constraint='y >= 1' invariant='5' (w=1, n_children=0, n_descendants=0)
-├── [1.2.1.1] constraint='y >= (-10)' invariant='4' (w=1, n_children=0, n_descendants=0)
-├── [1.2.3] constraint='y >= 1' invariant=None (w=2, n_children=2, n_descendants=2)
-│   ├── [1.2.3.1.1] constraint='x >= 1' invariant='2' (w=1, n_children=0, n_descendants=0)
-│   └── [1.2.3.1.2.1] constraint='not (x <= y)' invariant='1' (w=1, n_children=0, n_descendants=0)
-└── [1.2.3.4.1] constraint='y <= (-11)' invariant='3' (w=1, n_children=0, n_descendants=0)\
+├── [1] constraint='x >= 1' invariant=None (w=4, n_children=3, n_descendants=5)
+│   ├── [1.1] constraint='y >= 1' invariant=None (w=2, n_children=2, n_descendants=2)
+│   │   ├── [1.1.1] constraint='x <= y' invariant='2' (w=1, n_children=0, n_descendants=0)
+│   │   └── [1.1.2] constraint='x > y' invariant='1' (w=1, n_children=0, n_descendants=0)
+│   ├── [1.2] constraint='y <= (-11)' invariant='3' (w=1, n_children=0, n_descendants=0)
+│   └── [1.3.1] constraint='y >= (-10)' invariant='4' (w=1, n_children=0, n_descendants=0)
+├── [2.1] constraint='x <= 0' invariant='5' (w=1, n_children=0, n_descendants=0)
+└── [3.1] constraint='y <= 0' invariant='6' (w=1, n_children=0, n_descendants=0)\
 """)
     hdr_dict_hierarchy = hdr.to_dict_hierarchical()
     hdr_dict_flat = hdr.to_dict_flat()
@@ -128,128 +130,117 @@ def _collect_hierarchy(groups: list[RegionGroup]) -> list[tuple[str, int]]:
 - !RegionGroup
   label_path: '1'
   constraints:
-  - x <= 0
-  introduced_constraint: x <= 0
-  weight: 2
-  n_children_regions: 2
-  n_descendant_regions: 2
-  n_leaf_regions: 2
+  - x >= 1
+  introduced_constraint: x >= 1
+  weight: 4
+  n_children_regions: 3
+  n_descendant_regions: 5
+  n_leaf_regions: 4
   children:
   - !RegionGroup
     label_path: '1.1'
     constraints:
-    - x <= 0
-    - y <= 0
-    introduced_constraint: y <= 0
+    - x >= 1
+    - y >= 1
+    introduced_constraint: y >= 1
+    weight: 2
+    n_children_regions: 2
+    n_descendant_regions: 2
+    n_leaf_regions: 2
+    children:
+    - !RegionGroup
+      label_path: 1.1.1
+      constraints:
+      - x >= 1
+      - y >= 1
+      - x <= y
+      introduced_constraint: x <= y
+      weight: 1
+      n_children_regions: 0
+      n_descendant_regions: 0
+      n_leaf_regions: 1
+      invariant: '2'
+      example_input:
+        x: '1'
+        y: '1'
+      example_output: '2'
+    - !RegionGroup
+      label_path: 1.1.2
+      constraints:
+      - x >= 1
+      - y >= 1
+      - x > y
+      introduced_constraint: x > y
+      weight: 1
+      n_children_regions: 0
+      n_descendant_regions: 0
+      n_leaf_regions: 1
+      invariant: '1'
+      example_input:
+        x: '2'
+        y: '1'
+      example_output: '1'
+  - !RegionGroup
+    label_path: '1.2'
+    constraints:
+    - x >= 1
+    - y <= (-11)
+    introduced_constraint: y <= (-11)
     weight: 1
     n_children_regions: 0
     n_descendant_regions: 0
     n_leaf_regions: 1
-    invariant: '6'
+    invariant: '3'
     example_input:
-      x: '0'
-      y: '0'
-    example_output: '6'
+      x: '1'
+      y: (-11)
+    example_output: '3'
   - !RegionGroup
-    label_path: 1.1.2
+    label_path: 1.3.1
     constraints:
-    - x <= 0
+    - x >= 1
     - y <= 0
-    - y >= 1
-    introduced_constraint: y >= 1
+    - y >= (-10)
+    introduced_constraint: y >= (-10)
     weight: 1
     n_children_regions: 0
     n_descendant_regions: 0
     n_leaf_regions: 1
-    invariant: '5'
+    invariant: '4'
     example_input:
-      x: '0'
-      y: '1'
-    example_output: '5'
+      x: '1'
+      y: '0'
+    example_output: '4'
 - !RegionGroup
-  label_path: 1.2.1.1
+  label_path: '2.1'
   constraints:
+  - y >= 1
   - x <= 0
-  - y <= 0
-  - x >= 1
-  - y >= (-10)
-  introduced_constraint: y >= (-10)
+  introduced_constraint: x <= 0
   weight: 1
   n_children_regions: 0
   n_descendant_regions: 0
   n_leaf_regions: 1
-  invariant: '4'
+  invariant: '5'
   example_input:
-    x: '1'
-    y: '0'
-  example_output: '4'
-- !RegionGroup
-  label_path: 1.2.3
-  constraints:
-  - x <= 0
-  - y <= 0
-  - y >= 1
-  introduced_constraint: y >= 1
-  weight: 2
-  n_children_regions: 2
-  n_descendant_regions: 2
-  n_leaf_regions: 2
-  children:
-  - !RegionGroup
-    label_path: 1.2.3.1.1
-    constraints:
-    - x <= 0
-    - y <= 0
-    - y >= 1
-    - x <= y
-    - x >= 1
-    introduced_constraint: x >= 1
-    weight: 1
-    n_children_regions: 0
-    n_descendant_regions: 0
-    n_leaf_regions: 1
-    invariant: '2'
-    example_input:
-      x: '1'
-      y: '1'
-    example_output: '2'
-  - !RegionGroup
-    label_path: 1.2.3.1.2.1
-    constraints:
-    - x <= 0
-    - y <= 0
-    - y >= 1
-    - x <= y
-    - x >= 1
-    - not (x <= y)
-    introduced_constraint: not (x <= y)
-    weight: 1
-    n_children_regions: 0
-    n_descendant_regions: 0
-    n_leaf_regions: 1
-    invariant: '1'
-    example_input:
-      x: '2'
-      y: '1'
-    example_output: '1'
+    x: '0'
+    y: '1'
+  example_output: '5'
 - !RegionGroup
-  label_path: 1.2.3.4.1
+  label_path: '3.1'
   constraints:
   - x <= 0
   - y <= 0
-  - y >= 1
-  - x >= 1
-  - y <= (-11)
-  introduced_constraint: y <= (-11)
+  introduced_constraint: y <= 0
   weight: 1
   n_children_regions: 0
   n_descendant_regions: 0
   n_leaf_regions: 1
-  invariant: '3'
+  invariant: '6'
   example_input:
-    x: '1'
-    y: (-11)
-  example_output: '3'
+    x: '0'
+    y: '0'
+  example_output: '6'
 """)
 
     assert (hdr.dumper_func()(hdr_dict_flat)) == snapshot("""\
@@ -259,69 +250,21 @@ def _collect_hierarchy(groups: list[RegionGroup]) -> list[tuple[str, int]]:
 - id: '0'
   label_path: '1'
   depth: 1
-  weight: 2
+  weight: 4
   constraints:
-  - x <= 0
-  introduced_constraint: x <= 0
-  n_children_regions: 2
+  - x >= 1
+  introduced_constraint: x >= 1
+  n_children_regions: 3
   children:
+  - '3'
   - '4'
   - '5'
-- id: '4'
+- id: '3'
   label_path: '1.1'
   depth: 2
-  weight: 1
-  constraints:
-  - x <= 0
-  - y <= 0
-  introduced_constraint: y <= 0
-  invariant: '6'
-  example_input:
-    x: '0'
-    y: '0'
-  example_output: '6'
-  n_children_regions: 0
-  children: []
-- id: '5'
-  label_path: 1.1.2
-  depth: 2
-  weight: 1
-  constraints:
-  - x <= 0
-  - y <= 0
-  - y >= 1
-  introduced_constraint: y >= 1
-  invariant: '5'
-  example_input:
-    x: '0'
-    y: '1'
-  example_output: '5'
-  n_children_regions: 0
-  children: []
-- id: '1'
-  label_path: 1.2.1.1
-  depth: 1
-  weight: 1
-  constraints:
-  - x <= 0
-  - y <= 0
-  - x >= 1
-  - y >= (-10)
-  introduced_constraint: y >= (-10)
-  invariant: '4'
-  example_input:
-    x: '1'
-    y: '0'
-  example_output: '4'
-  n_children_regions: 0
-  children: []
-- id: '2'
-  label_path: 1.2.3
-  depth: 1
   weight: 2
   constraints:
-  - x <= 0
-  - y <= 0
+  - x >= 1
   - y >= 1
   introduced_constraint: y >= 1
   n_children_regions: 2
@@ -329,16 +272,14 @@ def _collect_hierarchy(groups: list[RegionGroup]) -> list[tuple[str, int]]:
   - '6'
   - '7'
 - id: '6'
-  label_path: 1.2.3.1.1
-  depth: 2
+  label_path: 1.1.1
+  depth: 3
   weight: 1
   constraints:
-  - x <= 0
-  - y <= 0
+  - x >= 1
   - y >= 1
   - x <= y
-  - x >= 1
-  introduced_constraint: x >= 1
+  introduced_constraint: x <= y
   invariant: '2'
   example_input:
     x: '1'
@@ -347,17 +288,14 @@ def _collect_hierarchy(groups: list[RegionGroup]) -> list[tuple[str, int]]:
   n_children_regions: 0
   children: []
 - id: '7'
-  label_path: 1.2.3.1.2.1
-  depth: 2
+  label_path: 1.1.2
+  depth: 3
   weight: 1
   constraints:
-  - x <= 0
-  - y <= 0
-  - y >= 1
-  - x <= y
   - x >= 1
-  - not (x <= y)
-  introduced_constraint: not (x <= y)
+  - y >= 1
+  - x > y
+  introduced_constraint: x > y
   invariant: '1'
   example_input:
     x: '2'
@@ -365,14 +303,11 @@ def _collect_hierarchy(groups: list[RegionGroup]) -> list[tuple[str, int]]:
   example_output: '1'
   n_children_regions: 0
   children: []
-- id: '3'
-  label_path: 1.2.3.4.1
-  depth: 1
+- id: '4'
+  label_path: '1.2'
+  depth: 2
   weight: 1
   constraints:
-  - x <= 0
-  - y <= 0
-  - y >= 1
   - x >= 1
   - y <= (-11)
   introduced_constraint: y <= (-11)
@@ -383,4 +318,50 @@ def _collect_hierarchy(groups: list[RegionGroup]) -> list[tuple[str, int]]:
   example_output: '3'
   n_children_regions: 0
   children: []
+- id: '5'
+  label_path: 1.3.1
+  depth: 2
+  weight: 1
+  constraints:
+  - x >= 1
+  - y <= 0
+  - y >= (-10)
+  introduced_constraint: y >= (-10)
+  invariant: '4'
+  example_input:
+    x: '1'
+    y: '0'
+  example_output: '4'
+  n_children_regions: 0
+  children: []
+- id: '1'
+  label_path: '2.1'
+  depth: 1
+  weight: 1
+  constraints:
+  - y >= 1
+  - x <= 0
+  introduced_constraint: x <= 0
+  invariant: '5'
+  example_input:
+    x: '0'
+    y: '1'
+  example_output: '5'
+  n_children_regions: 0
+  children: []
+- id: '2'
+  label_path: '3.1'
+  depth: 1
+  weight: 1
+  constraints:
+  - x <= 0
+  - y <= 0
+  introduced_constraint: y <= 0
+  invariant: '6'
+  example_input:
+    x: '0'
+    y: '0'
+  example_output: '6'
+  n_children_regions: 0
+  children: []
 """)

From fd4964f6766c9014e821784f89941d4ee449faa0 Mon Sep 17 00:00:00 2001
From: hongyu <linerwang1997@gmail.com>
Date: Mon, 29 Jun 2026 16:39:05 +0100
Subject: [PATCH 02/11] FEAT!: unify grouped decomp res and raw decomp res

---
 .../region_decomp/__init__.py                 | 361 ++++--------------
 .../tests/test_grouped_region_decomp.py       | 323 ++--------------
 2 files changed, 104 insertions(+), 580 deletions(-)

diff --git a/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py b/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
index 7e41a99d..4bd15986 100644
--- a/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
+++ b/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
@@ -1,142 +1,21 @@
-"""Post-processing (Hierarchical groupping) for region decomposition."""
+"""Post-processing (hierarchical grouping) for region decomposition."""
 
 from __future__ import annotations
 
-import itertools
 from collections.abc import Callable
-from dataclasses import dataclass
-from functools import partial, reduce
-from typing import Any, Literal, NoReturn, Protocol, Self, TypedDict
+from functools import reduce
+from typing import Any, NoReturn, Self, TypedDict
 
-import yaml
 from devtools import pformat
 from imandrax_api.lib import RegionStr
+from pydantic import BaseModel, model_validator
 
-from imandrax_api_models.proto_models import DecomposeRes, Error
-from imandrax_api_models.yaml_utils import str_representer
+from imandrax_api_models.proto_models import DecomposeRes
 
 from .icicle_widget import mk_icicle_widget_html
 
 
-@dataclass
-class HumDecomposeRes:
-    """Human readable decomp result"""
-
-    _variants: (
-        tuple[Literal['Success'], list[RegionGroup]]
-        | tuple[Literal['Fail'], list[Error]]
-    )
-
-    @classmethod
-    def mk_success(cls, pl: list[RegionGroup]) -> Self:
-        return cls(_variants=('Success', pl))
-
-    @classmethod
-    def mk_fail(cls, pl: list[Error]) -> Self:
-        return cls(_variants=('Fail', pl))
-
-    @classmethod
-    def from_decomp_res(cls, v: DecomposeRes) -> HumDecomposeRes:
-        return hum_of_decomp_res(v)
-
-    @classmethod
-    def from_regions(cls, regions: list[RegionStr]) -> HumDecomposeRes:
-        groups = group_regions(regions)
-        return HumDecomposeRes.mk_success(groups)
-
-    def to_tree_str(
-        self,
-        *,
-        depth_limit: int | None = None,
-        summarize: RegionGroupSummarizer | None = None,
-    ) -> str:
-        match self._variants:
-            case ('Fail', errs):
-                return pformat(errs, indent=2)
-            case ('Success', groups):
-                summarize_ = summarize or default_region_group_summary
-                lines: list[str] = []
-                for i, group in enumerate(groups):
-                    is_last = i == len(groups) - 1
-                    _tree_lines(
-                        lines,
-                        group,
-                        prefix='',
-                        is_last=is_last,
-                        depth_limit=depth_limit,
-                        summarize=summarize_,
-                    )
-                return '\n'.join(lines)
-
-    def summary(self) -> dict[str, str | int] | None:
-        match self._variants:
-            case ('Fail', _):
-                return None
-            case ('Success', groups):
-                return {
-                    'n_regions': _count_leaf_regions(groups),
-                    # 'n_regions': sum(rg.n_regions() for rg in groups),
-                    # 'n_leaf_regions': _count_leaf_regions(groups),
-                    # 'max_depth': _max_tree_depth_of_groups(groups),
-                    # 'max_label_depth': _max_label_depth_of_groups(groups),
-                    # 'max_constraints_per_region': _max_constraints_per_region(groups),
-                }
-
-    @staticmethod
-    def dumper_func(depth_limit: int | None = None) -> Callable[..., str]:
-        dumper_cls = mk_dumper(depth_limit=depth_limit)
-        return partial(
-            yaml.dump,
-            Dumper=dumper_cls,
-            default_flow_style=False,
-            sort_keys=False,
-        )
-
-    def to_dict_hierarchical(
-        self,
-        depth_limit: int | None = None,
-    ) -> dict[str, Any]:
-        match self._variants:
-            case ('Fail', errs):
-                return {'errors': errs}
-            case ('Success', groups):
-                return {'summary': self.summary(), 'region_groups': groups}
-
-    def to_dict_flat(
-        self,
-    ) -> dict[str, Any]:
-        match self._variants:
-            case ('Fail', errs):
-                return {'errors': errs}
-            case ('Success', groups):
-                indexed_groups = indexed_of_region_groups(groups)
-                return {'summary': self.summary(), 'region_groups': indexed_groups}
-
-    def _repr_html_(self) -> str:
-        match self._variants:
-            case ('Fail', errs):
-                return f'<pre>{pformat(errs, indent=2)}</pre>'
-            case ('Success', groups):
-                return mk_icicle_widget_html(groups)
-
-
-def hum_of_decomp_res(decomp_res: DecomposeRes) -> HumDecomposeRes:
-    if decomp_res.err is None:
-        regions = decomp_res.regions_str
-        if regions is None:
-            raise ValueError(
-                'DecomposeRes has no `regions_str`; humanizing regions requires '
-                'the decomposition to be requested with `string_results=True`.'
-            )
-        groups = group_regions(regions)
-        return HumDecomposeRes.mk_success(groups)
-
-    else:
-        return HumDecomposeRes.mk_fail(decomp_res.errors)
-
-
-@dataclass
-class RegionGroup:
+class RegionGroup(BaseModel):
     """
     A hierarchical group of regions sharing constraints.
 
@@ -160,8 +39,8 @@ class RegionGroup:
 
     constraints: list[str]
     label_path: list[int]
-    region: RegionStr | None
-    children: list[RegionGroup]
+    region: RegionStr | None = None
+    children: list[RegionGroup] = []
     weight: int
 
     def n_regions(self) -> int:
@@ -203,121 +82,86 @@ def to_json_dict(self) -> dict[str, Any]:
         return d
 
 
-def group_regions(regions: list[RegionStr]) -> list[RegionGroup]:
-    """Group regions hierarchically based on constraints."""
-    return _loop_group_regions([], [], regions)
-
-
-@dataclass
-class IndexedRegionGroup:
-    """Equivalent to RegionGroup, but in flat structure."""
-
-    id: str
-    constraints: list[str]
-    label_path: list[int]
-    region: RegionStr | None
-    children: list[str]
-    weight: int
-    depth: int  # 1-indexed depth
-
-
-def indexed_of_region_groups(groups: list[RegionGroup]) -> list[IndexedRegionGroup]:
-    id_generator = (str(i) for i in itertools.count())
-
-    def gen_id() -> str:
-        return next(id_generator)
-
-    def loop(
-        group_and_id_lst: list[tuple[RegionGroup, str]],
-        acc: list[IndexedRegionGroup],
-        depth: int,
-    ) -> None:
-        for rg, id in group_and_id_lst:
-            c_ids = [gen_id() for _ in rg.children]
-
-            irg = IndexedRegionGroup(
-                id=id,
-                constraints=rg.constraints,
-                label_path=rg.label_path,
-                region=rg.region,
-                children=c_ids,
-                weight=rg.weight,
-                depth=depth,
-            )
-            acc.append(irg)
-
-            children_group_and_id_lst = list(zip(rg.children, c_ids, strict=True))
-            loop(children_group_and_id_lst, acc, depth + 1)
-
-    initial_ids = [gen_id() for _ in groups]
-    group_and_id_lst = list(zip(groups, initial_ids, strict=True))
-    acc: list[IndexedRegionGroup] = []
-    loop(group_and_id_lst, acc, 1)
-    return acc
-
-
-def _max_label_depth_of_groups(groups: list[RegionGroup]) -> int:
+class HumDecomposeRes(DecomposeRes):
     """
-    Max length of any label_path in the tree (index-chain depth).
+    A `DecomposeRes` augmented with hierarchical region grouping.
 
-    Counts collapsed singleton levels since label_path is built before
-    singleton promotion in `_loop_group_regions`.
+    `region_groups` is derived from the inherited `regions_str` (a pure function
+    of it), and is auto-populated on validation. Failure is represented by the
+    inherited `err` / `errors`, exactly as on `DecomposeRes`.
     """
-    init: int = 1
-
-    def loop(groups: list[RegionGroup]):
-        label_lengths: list[int] = [len(rg.label_path) for rg in groups]
-        nonlocal init
-        init = max(init, (max(label_lengths) if label_lengths else 0))
-        for rg in groups:
-            loop(rg.children)
-
-    loop(groups)
-    return init
 
+    region_groups: list[RegionGroup] = []
 
-def _max_tree_depth_of_groups(groups: list[RegionGroup]) -> int:
-    """Max nesting depth of RegionGroup nodes (matches visual tree)."""
+    @model_validator(mode='after')
+    def _populate_region_groups(self) -> Self:
+        if not self.region_groups and self.regions_str:
+            self.region_groups = group_regions(self.regions_str)
+        return self
 
-    def depth(rg: RegionGroup) -> int:
-        return 1 + max((depth(c) for c in rg.children), default=0)
-
-    return max((depth(g) for g in groups), default=0)
+    @classmethod
+    def from_decomp_res(cls, v: DecomposeRes) -> HumDecomposeRes:
+        return cls.model_validate(v.model_dump())
 
+    def to_tree_str(
+        self,
+        *,
+        depth_limit: int | None = None,
+        summarize: Callable[[RegionGroup], str] | None = None,
+    ) -> str:
+        if self.errors:
+            return pformat(self.errors, indent=2)
+        return render_region_groups(
+            self.region_groups, depth_limit=depth_limit, summarize=summarize
+        )
 
-def _count_leaf_regions(groups: list[RegionGroup]) -> int:
-    """Count of nodes carrying a concrete RegionStr (region is not None)."""
-    total = 0
-    for g in groups:
-        if g.region is not None:
-            total += 1
-        total += _count_leaf_regions(g.children)
-    return total
+    def _repr_html_(self) -> str:
+        if self.errors:
+            return f'<pre>{pformat(self.errors, indent=2)}</pre>'
+        return mk_icicle_widget_html(self.region_groups)
 
 
-def _max_constraints_per_region(groups: list[RegionGroup]) -> int:
-    """Longest constraints list across all nodes."""
+def group_regions(regions: list[RegionStr]) -> list[RegionGroup]:
+    """Group regions hierarchically based on constraints."""
+    return _loop_group_regions([], [], regions)
 
-    def walk(rg: RegionGroup) -> int:
-        return max(len(rg.constraints), *(walk(c) for c in rg.children), 0)
 
-    return max((walk(g) for g in groups), default=0)
+# Tree rendering
+# ====================
 
 
-class RegionGroupSummarizer(Protocol):
-    def __call__(self, group: RegionGroup) -> str: ...
+def render_region_groups(
+    groups: list[RegionGroup],
+    *,
+    depth_limit: int | None = None,
+    summarize: Callable[[RegionGroup], str] | None = None,
+) -> str:
+    """Render a forest of `RegionGroup`s as an ASCII tree."""
+    summarize_ = summarize or default_region_group_summary
+    lines: list[str] = []
+    for i, group in enumerate(groups):
+        is_last = i == len(groups) - 1
+        _tree_lines(
+            lines,
+            group,
+            prefix='',
+            is_last=is_last,
+            depth_limit=depth_limit,
+            summarize=summarize_,
+        )
+    return '\n'.join(lines)
 
 
 def default_region_group_summary(group: RegionGroup) -> str:
     label = '.'.join(map(str, group.label_path))
-    # rg_constraints is the full path from root; [-1] is this node's own constraint.
+    # constraints is the full path from root; [-1] is this node's own constraint.
     constraint = group.constraints[-1] if group.constraints else '?'
     invariant: str | None = None
     if (region := group.region) is not None:
         invariant = region.invariant_str
     parts = [
         f'[{label}]',
-        f'{constraint=}',
+        f'constraints[-1]={constraint}',
         f'{invariant=}',
         f'(w={group.weight}, n_children={len(group.children)}, n_descendants={group.n_descendant_regions()})',
     ]
@@ -331,7 +175,7 @@ def _tree_lines(
     prefix: str,
     is_last: bool,
     depth_limit: int | None,
-    summarize: RegionGroupSummarizer,
+    summarize: Callable[[RegionGroup], str],
 ) -> None:
     connector = '└── ' if is_last else '├── '
     lines.append(f'{prefix}{connector}{summarize(group)}')
@@ -354,6 +198,10 @@ def _tree_lines(
         )
 
 
+# Grouping algorithm
+# ====================
+
+
 def _loop_group_regions(
     idx_path: list[int], constraint_path: list[str], regions: list[RegionStr]
 ) -> list[RegionGroup]:
@@ -509,74 +357,3 @@ def loop(
         groups=[], regions=regions, idx_path=idx_path, constraint_path=constraint_path
     )
     return reduce(loop, constraints_by_most_frequent, init)['groups'][::-1]
-
-
-# YAML Dump
-# ====================
-
-
-def _region_str_representer(dumper: yaml.Dumper, data: RegionStr) -> yaml.Node:
-    mapping: dict[str, object] = {}
-    if data.constraints_str is not None:
-        mapping['constraints'] = data.constraints_str
-    if data.invariant_str is not None:
-        mapping['invariant'] = data.invariant_str
-    if data.model_str is not None:
-        mapping['model'] = data.model_str
-    if data.model_eval_str is not None:
-        mapping['model_eval'] = data.model_eval_str
-    return dumper.represent_mapping('!Region', mapping)
-
-
-def _region_group_representer(
-    dumper: yaml.Dumper, data: RegionGroup, *, depth_limit: int | None = None
-) -> yaml.Node:
-    mapping: dict[str, object] = data.to_dict()
-    if data.children:
-        if depth_limit is None or depth_limit <= 0:
-            mapping['children'] = data.children
-    return dumper.represent_mapping('!RegionGroup', mapping)
-
-
-def _indexed_region_group_representer(
-    dumper: yaml.Dumper, data: IndexedRegionGroup
-) -> yaml.Node:
-    mapping: dict[str, object] = {}
-
-    mapping['id'] = data.id
-    mapping['label_path'] = '.'.join(map(str, data.label_path))
-    mapping['depth'] = data.depth
-    mapping['weight'] = data.weight
-
-    mapping['constraints'] = data.constraints
-    mapping['introduced_constraint'] = data.constraints[-1]
-    if (region := data.region) is not None:
-        mapping['invariant'] = region.invariant_str
-        mapping['example_input'] = region.model_str
-        mapping['example_output'] = region.model_eval_str
-
-    mapping['n_children_regions'] = len(data.children)
-    mapping['children'] = data.children
-    return dumper.represent_mapping('tag:yaml.org,2002:map', mapping)
-    # return dumper.represent_mapping('!IndexedRegionGroup', mapping)
-
-
-def mk_dumper(*, depth_limit: int | None = None) -> type[yaml.Dumper]:
-    class RegionDecompDumper(yaml.Dumper):
-        pass
-
-    RegionDecompDumper.add_representer(
-        RegionStr,
-        _region_str_representer,
-    )
-
-    def _rg_representer(dumper: yaml.Dumper, data: RegionGroup) -> yaml.Node:
-        next_limit = None if depth_limit is None else depth_limit - 1
-        return _region_group_representer(dumper, data, depth_limit=next_limit)
-
-    RegionDecompDumper.add_representer(RegionGroup, _rg_representer)
-    RegionDecompDumper.add_representer(
-        IndexedRegionGroup, _indexed_region_group_representer
-    )
-    RegionDecompDumper.add_representer(str, str_representer)
-    return RegionDecompDumper
diff --git a/packages/imandrax-api-models/tests/test_grouped_region_decomp.py b/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
index 7c8a29cd..3244c1d0 100644
--- a/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
+++ b/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
@@ -4,12 +4,11 @@
 from imandrax_api.lib import RegionStr
 from inline_snapshot import snapshot
 
+from imandrax_api_models.proto_models import DecomposeRes
 from imandrax_api_models.region_decomp import HumDecomposeRes, RegionGroup
 
 
-def trust():
-    from typing import NoReturn
-
+def trust() -> DecomposeRes:
     import dotenv
 
     from imandrax_api_models.client import ImandraXClient
@@ -32,22 +31,15 @@ def trust():
         if y > 0 then 5
         else 6"""
     _eval_res = c.eval_src(IML_CODE)
-    decomp_res = c.decompose(name='classify', string_results=True, prune=True)
-
-    def raise_(exc: BaseException) -> NoReturn:
-        raise exc
-
-    regions: list[RegionStr] = (
-        decomp_res.regions_str
-        if (decomp_res.regions_str)
-        else (raise_(ValueError('No regions')))
-    )
-    return regions
+    return c.decompose(name='classify', string_results=True, prune=True)
 
 
 def test():
-    regions = trust()
-    assert regions == snapshot(
+    decomp_res = trust()
+    hdr = HumDecomposeRes.from_decomp_res(decomp_res)
+
+    # The raw regions are inherited from DecomposeRes.
+    assert hdr.regions_str == snapshot(
         [
             RegionStr(
                 constraints_str=['y <= 0', 'x <= 0'],
@@ -88,280 +80,35 @@ def test():
         ]
     )
 
-    hdr = HumDecomposeRes.from_regions(regions)
+    # region_groups is auto-populated on validation from regions_str.
     assert hdr.to_tree_str() == snapshot("""\
-├── [1] constraint='x >= 1' invariant=None (w=4, n_children=3, n_descendants=5)
-│   ├── [1.1] constraint='y >= 1' invariant=None (w=2, n_children=2, n_descendants=2)
-│   │   ├── [1.1.1] constraint='x <= y' invariant='2' (w=1, n_children=0, n_descendants=0)
-│   │   └── [1.1.2] constraint='x > y' invariant='1' (w=1, n_children=0, n_descendants=0)
-│   ├── [1.2] constraint='y <= (-11)' invariant='3' (w=1, n_children=0, n_descendants=0)
-│   └── [1.3.1] constraint='y >= (-10)' invariant='4' (w=1, n_children=0, n_descendants=0)
-├── [2.1] constraint='x <= 0' invariant='5' (w=1, n_children=0, n_descendants=0)
-└── [3.1] constraint='y <= 0' invariant='6' (w=1, n_children=0, n_descendants=0)\
+├── [1] constraints[-1]=x >= 1 invariant=None (w=4, n_children=3, n_descendants=5)
+│   ├── [1.1] constraints[-1]=y >= 1 invariant=None (w=2, n_children=2, n_descendants=2)
+│   │   ├── [1.1.1] constraints[-1]=x <= y invariant='2' (w=1, n_children=0, n_descendants=0)
+│   │   └── [1.1.2] constraints[-1]=x > y invariant='1' (w=1, n_children=0, n_descendants=0)
+│   ├── [1.2] constraints[-1]=y <= (-11) invariant='3' (w=1, n_children=0, n_descendants=0)
+│   └── [1.3.1] constraints[-1]=y >= (-10) invariant='4' (w=1, n_children=0, n_descendants=0)
+├── [2.1] constraints[-1]=x <= 0 invariant='5' (w=1, n_children=0, n_descendants=0)
+└── [3.1] constraints[-1]=y <= 0 invariant='6' (w=1, n_children=0, n_descendants=0)\
 """)
-    hdr_dict_hierarchy = hdr.to_dict_hierarchical()
-    hdr_dict_flat = hdr.to_dict_flat()
-    label_path_n_children_map_flat: list[tuple[str, int]] = [
-        ('.'.join(map(str, irg.label_path)), len(irg.children))
-        for irg in hdr_dict_flat['region_groups']
-    ]
 
-    # ::: test-child-count
-    def _collect_hierarchy(groups: list[RegionGroup]) -> list[tuple[str, int]]:
-        result: list[tuple[str, int]] = []
-        for rg in groups:
-            result.append(('.'.join(map(str, rg.label_path)), len(rg.children)))
-            if rg.children:
-                result.extend(_collect_hierarchy(rg.children))
-        return result
+    # (label_path, full constraint path) for every node, depth-first.
+    def _walk(groups: list[RegionGroup]) -> list[tuple[str, list[str]]]:
+        out: list[tuple[str, list[str]]] = []
+        for g in groups:
+            out.append(('.'.join(map(str, g.label_path)), g.constraints))
+            out.extend(_walk(g.children))
+        return out
 
-    label_path_n_children_map_hierarchy: list[tuple[str, int]] = _collect_hierarchy(
-        hdr_dict_hierarchy['region_groups']
-    )
-    assert sorted(label_path_n_children_map_flat) == sorted(
-        label_path_n_children_map_hierarchy
+    assert _walk(hdr.region_groups) == snapshot(
+        [
+            ('1', ['x >= 1']),
+            ('1.1', ['x >= 1', 'y >= 1']),
+            ('1.1.1', ['x >= 1', 'y >= 1', 'x <= y']),
+            ('1.1.2', ['x >= 1', 'y >= 1', 'x > y']),
+            ('1.2', ['x >= 1', 'y <= (-11)']),
+            ('1.3.1', ['x >= 1', 'y <= 0', 'y >= (-10)']),
+            ('2.1', ['y >= 1', 'x <= 0']),
+            ('3.1', ['x <= 0', 'y <= 0']),
+        ]
     )
-    # :::
-
-    assert (hdr.dumper_func()(hdr_dict_hierarchy)) == snapshot("""\
-summary:
-  n_regions: 6
-region_groups:
-- !RegionGroup
-  label_path: '1'
-  constraints:
-  - x >= 1
-  introduced_constraint: x >= 1
-  weight: 4
-  n_children_regions: 3
-  n_descendant_regions: 5
-  n_leaf_regions: 4
-  children:
-  - !RegionGroup
-    label_path: '1.1'
-    constraints:
-    - x >= 1
-    - y >= 1
-    introduced_constraint: y >= 1
-    weight: 2
-    n_children_regions: 2
-    n_descendant_regions: 2
-    n_leaf_regions: 2
-    children:
-    - !RegionGroup
-      label_path: 1.1.1
-      constraints:
-      - x >= 1
-      - y >= 1
-      - x <= y
-      introduced_constraint: x <= y
-      weight: 1
-      n_children_regions: 0
-      n_descendant_regions: 0
-      n_leaf_regions: 1
-      invariant: '2'
-      example_input:
-        x: '1'
-        y: '1'
-      example_output: '2'
-    - !RegionGroup
-      label_path: 1.1.2
-      constraints:
-      - x >= 1
-      - y >= 1
-      - x > y
-      introduced_constraint: x > y
-      weight: 1
-      n_children_regions: 0
-      n_descendant_regions: 0
-      n_leaf_regions: 1
-      invariant: '1'
-      example_input:
-        x: '2'
-        y: '1'
-      example_output: '1'
-  - !RegionGroup
-    label_path: '1.2'
-    constraints:
-    - x >= 1
-    - y <= (-11)
-    introduced_constraint: y <= (-11)
-    weight: 1
-    n_children_regions: 0
-    n_descendant_regions: 0
-    n_leaf_regions: 1
-    invariant: '3'
-    example_input:
-      x: '1'
-      y: (-11)
-    example_output: '3'
-  - !RegionGroup
-    label_path: 1.3.1
-    constraints:
-    - x >= 1
-    - y <= 0
-    - y >= (-10)
-    introduced_constraint: y >= (-10)
-    weight: 1
-    n_children_regions: 0
-    n_descendant_regions: 0
-    n_leaf_regions: 1
-    invariant: '4'
-    example_input:
-      x: '1'
-      y: '0'
-    example_output: '4'
-- !RegionGroup
-  label_path: '2.1'
-  constraints:
-  - y >= 1
-  - x <= 0
-  introduced_constraint: x <= 0
-  weight: 1
-  n_children_regions: 0
-  n_descendant_regions: 0
-  n_leaf_regions: 1
-  invariant: '5'
-  example_input:
-    x: '0'
-    y: '1'
-  example_output: '5'
-- !RegionGroup
-  label_path: '3.1'
-  constraints:
-  - x <= 0
-  - y <= 0
-  introduced_constraint: y <= 0
-  weight: 1
-  n_children_regions: 0
-  n_descendant_regions: 0
-  n_leaf_regions: 1
-  invariant: '6'
-  example_input:
-    x: '0'
-    y: '0'
-  example_output: '6'
-""")
-
-    assert (hdr.dumper_func()(hdr_dict_flat)) == snapshot("""\
-summary:
-  n_regions: 6
-region_groups:
-- id: '0'
-  label_path: '1'
-  depth: 1
-  weight: 4
-  constraints:
-  - x >= 1
-  introduced_constraint: x >= 1
-  n_children_regions: 3
-  children:
-  - '3'
-  - '4'
-  - '5'
-- id: '3'
-  label_path: '1.1'
-  depth: 2
-  weight: 2
-  constraints:
-  - x >= 1
-  - y >= 1
-  introduced_constraint: y >= 1
-  n_children_regions: 2
-  children:
-  - '6'
-  - '7'
-- id: '6'
-  label_path: 1.1.1
-  depth: 3
-  weight: 1
-  constraints:
-  - x >= 1
-  - y >= 1
-  - x <= y
-  introduced_constraint: x <= y
-  invariant: '2'
-  example_input:
-    x: '1'
-    y: '1'
-  example_output: '2'
-  n_children_regions: 0
-  children: []
-- id: '7'
-  label_path: 1.1.2
-  depth: 3
-  weight: 1
-  constraints:
-  - x >= 1
-  - y >= 1
-  - x > y
-  introduced_constraint: x > y
-  invariant: '1'
-  example_input:
-    x: '2'
-    y: '1'
-  example_output: '1'
-  n_children_regions: 0
-  children: []
-- id: '4'
-  label_path: '1.2'
-  depth: 2
-  weight: 1
-  constraints:
-  - x >= 1
-  - y <= (-11)
-  introduced_constraint: y <= (-11)
-  invariant: '3'
-  example_input:
-    x: '1'
-    y: (-11)
-  example_output: '3'
-  n_children_regions: 0
-  children: []
-- id: '5'
-  label_path: 1.3.1
-  depth: 2
-  weight: 1
-  constraints:
-  - x >= 1
-  - y <= 0
-  - y >= (-10)
-  introduced_constraint: y >= (-10)
-  invariant: '4'
-  example_input:
-    x: '1'
-    y: '0'
-  example_output: '4'
-  n_children_regions: 0
-  children: []
-- id: '1'
-  label_path: '2.1'
-  depth: 1
-  weight: 1
-  constraints:
-  - y >= 1
-  - x <= 0
-  introduced_constraint: x <= 0
-  invariant: '5'
-  example_input:
-    x: '0'
-    y: '1'
-  example_output: '5'
-  n_children_regions: 0
-  children: []
-- id: '2'
-  label_path: '3.1'
-  depth: 1
-  weight: 1
-  constraints:
-  - x <= 0
-  - y <= 0
-  introduced_constraint: y <= 0
-  invariant: '6'
-  example_input:
-    x: '0'
-    y: '0'
-  example_output: '6'
-  n_children_regions: 0
-  children: []
-""")

From 31e630378e546f620b7c595d36d98950aa00e458 Mon Sep 17 00:00:00 2001
From: hongyu <linerwang1997@gmail.com>
Date: Mon, 29 Jun 2026 16:46:27 +0100
Subject: [PATCH 03/11] REFA: renames

---
 .../region_decomp/__init__.py                 | 89 +++++++++----------
 1 file changed, 43 insertions(+), 46 deletions(-)

diff --git a/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py b/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
index 4bd15986..6c312bfd 100644
--- a/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
+++ b/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
@@ -8,13 +8,49 @@
 
 from devtools import pformat
 from imandrax_api.lib import RegionStr
-from pydantic import BaseModel, model_validator
+from pydantic import BaseModel, Field, model_validator
 
 from imandrax_api_models.proto_models import DecomposeRes
 
 from .icicle_widget import mk_icicle_widget_html
 
 
+class HumDecomposeRes(DecomposeRes):
+    """A `DecomposeRes` augmented with hierarchical region grouping."""
+
+    region_groups: list[RegionGroup] = Field(
+        default_factory=list,
+        description='Region groups grouped by constraints, containing child groups recursively. Empty when no regions are available (decomposition error).',
+    )
+
+    @model_validator(mode='after')
+    def _populate_region_groups(self) -> Self:
+        if not self.region_groups and self.regions_str:
+            self.region_groups = group_regions(self.regions_str)
+        return self
+
+    @classmethod
+    def from_decomp_res(cls, v: DecomposeRes) -> HumDecomposeRes:
+        return cls.model_validate(v.model_dump())
+
+    def to_tree_str(
+        self,
+        *,
+        depth_limit: int | None = None,
+        summarize: Callable[[RegionGroup], str] | None = None,
+    ) -> str:
+        if self.errors:
+            return pformat(self.errors, indent=2)
+        return render_region_groups(
+            self.region_groups, depth_limit=depth_limit, tree_repr=summarize
+        )
+
+    def _repr_html_(self) -> str:
+        if self.errors:
+            return f'<pre>{pformat(self.errors, indent=2)}</pre>'
+        return mk_icicle_widget_html(self.region_groups)
+
+
 class RegionGroup(BaseModel):
     """
     A hierarchical group of regions sharing constraints.
@@ -82,45 +118,6 @@ def to_json_dict(self) -> dict[str, Any]:
         return d
 
 
-class HumDecomposeRes(DecomposeRes):
-    """
-    A `DecomposeRes` augmented with hierarchical region grouping.
-
-    `region_groups` is derived from the inherited `regions_str` (a pure function
-    of it), and is auto-populated on validation. Failure is represented by the
-    inherited `err` / `errors`, exactly as on `DecomposeRes`.
-    """
-
-    region_groups: list[RegionGroup] = []
-
-    @model_validator(mode='after')
-    def _populate_region_groups(self) -> Self:
-        if not self.region_groups and self.regions_str:
-            self.region_groups = group_regions(self.regions_str)
-        return self
-
-    @classmethod
-    def from_decomp_res(cls, v: DecomposeRes) -> HumDecomposeRes:
-        return cls.model_validate(v.model_dump())
-
-    def to_tree_str(
-        self,
-        *,
-        depth_limit: int | None = None,
-        summarize: Callable[[RegionGroup], str] | None = None,
-    ) -> str:
-        if self.errors:
-            return pformat(self.errors, indent=2)
-        return render_region_groups(
-            self.region_groups, depth_limit=depth_limit, summarize=summarize
-        )
-
-    def _repr_html_(self) -> str:
-        if self.errors:
-            return f'<pre>{pformat(self.errors, indent=2)}</pre>'
-        return mk_icicle_widget_html(self.region_groups)
-
-
 def group_regions(regions: list[RegionStr]) -> list[RegionGroup]:
     """Group regions hierarchically based on constraints."""
     return _loop_group_regions([], [], regions)
@@ -134,10 +131,10 @@ def render_region_groups(
     groups: list[RegionGroup],
     *,
     depth_limit: int | None = None,
-    summarize: Callable[[RegionGroup], str] | None = None,
+    tree_repr: Callable[[RegionGroup], str] | None = None,
 ) -> str:
-    """Render a forest of `RegionGroup`s as an ASCII tree."""
-    summarize_ = summarize or default_region_group_summary
+    """Render a forest of `RegionGroup`s as a tree in text."""
+    tree_repr_ = tree_repr or default_region_group_repr
     lines: list[str] = []
     for i, group in enumerate(groups):
         is_last = i == len(groups) - 1
@@ -147,12 +144,12 @@ def render_region_groups(
             prefix='',
             is_last=is_last,
             depth_limit=depth_limit,
-            summarize=summarize_,
+            summarize=tree_repr_,
         )
     return '\n'.join(lines)
 
 
-def default_region_group_summary(group: RegionGroup) -> str:
+def default_region_group_repr(group: RegionGroup) -> str:
     label = '.'.join(map(str, group.label_path))
     # constraints is the full path from root; [-1] is this node's own constraint.
     constraint = group.constraints[-1] if group.constraints else '?'
@@ -344,7 +341,7 @@ def loop(
         # `idx_path` / `constraint_path` describe this level and must stay
         # constant across reduce iterations. Only the `has` branch (recursed
         # above) gets the extended `new_idx_path` / `new_constraint_path`; the
-        # `without` regions handled by later iterations do NOT contain
+        # `without` regions handled by later iterations do not contain
         # `konstraint`, so it must not leak into their path.
         return Acc(
             groups=res[0],

From eddd8f90e1bdc07cecab87dfa9c9d914ec776732 Mon Sep 17 00:00:00 2001
From: hongyu <linerwang1997@gmail.com>
Date: Mon, 29 Jun 2026 16:47:23 +0100
Subject: [PATCH 04/11] REFA!(decomp): rename HumDecomposeRes to
 EnrichedDecomposeRes

---
 .../src/imandrax_api_models/region_decomp/__init__.py         | 4 ++--
 .../imandrax-api-models/tests/test_grouped_region_decomp.py   | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py b/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
index 6c312bfd..3a95dbe8 100644
--- a/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
+++ b/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
@@ -15,7 +15,7 @@
 from .icicle_widget import mk_icicle_widget_html
 
 
-class HumDecomposeRes(DecomposeRes):
+class EnrichedDecomposeRes(DecomposeRes):
     """A `DecomposeRes` augmented with hierarchical region grouping."""
 
     region_groups: list[RegionGroup] = Field(
@@ -30,7 +30,7 @@ def _populate_region_groups(self) -> Self:
         return self
 
     @classmethod
-    def from_decomp_res(cls, v: DecomposeRes) -> HumDecomposeRes:
+    def from_decomp_res(cls, v: DecomposeRes) -> EnrichedDecomposeRes:
         return cls.model_validate(v.model_dump())
 
     def to_tree_str(
diff --git a/packages/imandrax-api-models/tests/test_grouped_region_decomp.py b/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
index 3244c1d0..e2b1add0 100644
--- a/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
+++ b/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
@@ -5,7 +5,7 @@
 from inline_snapshot import snapshot
 
 from imandrax_api_models.proto_models import DecomposeRes
-from imandrax_api_models.region_decomp import HumDecomposeRes, RegionGroup
+from imandrax_api_models.region_decomp import EnrichedDecomposeRes, RegionGroup
 
 
 def trust() -> DecomposeRes:
@@ -36,7 +36,7 @@ def trust() -> DecomposeRes:
 
 def test():
     decomp_res = trust()
-    hdr = HumDecomposeRes.from_decomp_res(decomp_res)
+    hdr = EnrichedDecomposeRes.from_decomp_res(decomp_res)
 
     # The raw regions are inherited from DecomposeRes.
     assert hdr.regions_str == snapshot(

From 1c9053563ddeb87845e242595007c10f1c5b28a7 Mon Sep 17 00:00:00 2001
From: hongyu <linerwang1997@gmail.com>
Date: Mon, 29 Jun 2026 17:19:14 +0100
Subject: [PATCH 05/11] REFA: unify region group stat summary

---
 .../region_decomp/__init__.py                 | 67 +++++++++----------
 .../tests/test_grouped_region_decomp.py       | 16 ++---
 2 files changed, 39 insertions(+), 44 deletions(-)

diff --git a/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py b/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
index 3a95dbe8..335743dd 100644
--- a/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
+++ b/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
@@ -4,7 +4,7 @@
 
 from collections.abc import Callable
 from functools import reduce
-from typing import Any, NoReturn, Self, TypedDict
+from typing import NoReturn, Self, TypedDict
 
 from devtools import pformat
 from imandrax_api.lib import RegionStr
@@ -51,6 +51,11 @@ def _repr_html_(self) -> str:
         return mk_icicle_widget_html(self.region_groups)
 
 
+type JSONValue = str | int | float | bool | None | JSONObject | JSONArray
+type JSONObject = dict[str, JSONValue]
+type JSONArray = list[JSONValue]
+
+
 class RegionGroup(BaseModel):
     """
     A hierarchical group of regions sharing constraints.
@@ -88,34 +93,40 @@ def n_descendant_regions(self) -> int:
         return self.n_regions() - 1
 
     def n_leaf_regions(self) -> int:
-        """Total leaf regions in this subtree, including self if no children."""
+        """Total leaf regions in this subtree, counting self if no children."""
         if not self.children:
             return 1
         return sum(c.n_leaf_regions() for c in self.children)
 
-    def to_dict(self) -> dict[str, Any]:
-        """Core fields shared by all serialization formats."""
-        d: dict[str, Any] = {
-            'label_path': '.'.join(map(str, self.label_path)),
-            'constraints': self.constraints,
-            'introduced_constraint': self.constraints[-1] if self.constraints else '',
-            'weight': self.weight,
-            'n_children_regions': len(self.children),
-            'n_descendant_regions': self.n_descendant_regions(),
-            'n_leaf_regions': self.n_leaf_regions(),
-        }
+    def describe(self) -> JSONObject:
+        d: JSONObject = {}
+        d['label_path'] = '.'.join(map(str, self.label_path))
+        d['constraints'] = self.constraints
+        d['introduced_constraint'] = self.constraints[-1] if self.constraints else ''
+        d['weight'] = self.weight
+        d['n_children_regions'] = len(self.children)
+        d['n_descendant_regions'] = self.n_descendant_regions()
+        d['n_leaf_regions'] = self.n_leaf_regions()
         if (r := self.region) is not None:
             d['invariant'] = r.invariant_str
             d['example_input'] = r.model_str
             d['example_output'] = r.model_eval_str
         return d
 
-    def to_json_dict(self) -> dict[str, Any]:
-        """Serialize to a d3-hierarchy-compatible dict."""
-        d = self.to_dict()
-        if self.children:
-            d['children'] = [c.to_json_dict() for c in self.children]
-        return d
+    def repr_line(self) -> str:
+        """One-line representation of the region group."""
+        d = self.describe()
+        parts: list[str] = []
+        parts.append(f'[{d["label_path"]}]')
+        parts.append(f"new_constraint='{d['introduced_constraint']}'")
+        if d.get('invariant'):
+            parts.append(f"invariant='{d['invariant']}'")
+        n_leaf_regions = d['n_leaf_regions']
+        if n_leaf_regions != 1:
+            parts.append(f'n_leaf_regions={n_leaf_regions}')
+        else:
+            parts.append('is_leaf=True')
+        return ' '.join(parts)
 
 
 def group_regions(regions: list[RegionStr]) -> list[RegionGroup]:
@@ -134,7 +145,7 @@ def render_region_groups(
     tree_repr: Callable[[RegionGroup], str] | None = None,
 ) -> str:
     """Render a forest of `RegionGroup`s as a tree in text."""
-    tree_repr_ = tree_repr or default_region_group_repr
+    tree_repr_ = tree_repr or RegionGroup.repr_line
     lines: list[str] = []
     for i, group in enumerate(groups):
         is_last = i == len(groups) - 1
@@ -149,22 +160,6 @@ def render_region_groups(
     return '\n'.join(lines)
 
 
-def default_region_group_repr(group: RegionGroup) -> str:
-    label = '.'.join(map(str, group.label_path))
-    # constraints is the full path from root; [-1] is this node's own constraint.
-    constraint = group.constraints[-1] if group.constraints else '?'
-    invariant: str | None = None
-    if (region := group.region) is not None:
-        invariant = region.invariant_str
-    parts = [
-        f'[{label}]',
-        f'constraints[-1]={constraint}',
-        f'{invariant=}',
-        f'(w={group.weight}, n_children={len(group.children)}, n_descendants={group.n_descendant_regions()})',
-    ]
-    return ' '.join(parts)
-
-
 def _tree_lines(
     lines: list[str],
     group: RegionGroup,
diff --git a/packages/imandrax-api-models/tests/test_grouped_region_decomp.py b/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
index e2b1add0..7b786167 100644
--- a/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
+++ b/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
@@ -82,14 +82,14 @@ def test():
 
     # region_groups is auto-populated on validation from regions_str.
     assert hdr.to_tree_str() == snapshot("""\
-├── [1] constraints[-1]=x >= 1 invariant=None (w=4, n_children=3, n_descendants=5)
-│   ├── [1.1] constraints[-1]=y >= 1 invariant=None (w=2, n_children=2, n_descendants=2)
-│   │   ├── [1.1.1] constraints[-1]=x <= y invariant='2' (w=1, n_children=0, n_descendants=0)
-│   │   └── [1.1.2] constraints[-1]=x > y invariant='1' (w=1, n_children=0, n_descendants=0)
-│   ├── [1.2] constraints[-1]=y <= (-11) invariant='3' (w=1, n_children=0, n_descendants=0)
-│   └── [1.3.1] constraints[-1]=y >= (-10) invariant='4' (w=1, n_children=0, n_descendants=0)
-├── [2.1] constraints[-1]=x <= 0 invariant='5' (w=1, n_children=0, n_descendants=0)
-└── [3.1] constraints[-1]=y <= 0 invariant='6' (w=1, n_children=0, n_descendants=0)\
+├── [1] new_constraint='x >= 1' n_leaf_regions=4
+│   ├── [1.1] new_constraint='y >= 1' n_leaf_regions=2
+│   │   ├── [1.1.1] new_constraint='x <= y' invariant='2' is_leaf=True
+│   │   └── [1.1.2] new_constraint='x > y' invariant='1' is_leaf=True
+│   ├── [1.2] new_constraint='y <= (-11)' invariant='3' is_leaf=True
+│   └── [1.3.1] new_constraint='y >= (-10)' invariant='4' is_leaf=True
+├── [2.1] new_constraint='x <= 0' invariant='5' is_leaf=True
+└── [3.1] new_constraint='y <= 0' invariant='6' is_leaf=True\
 """)
 
     # (label_path, full constraint path) for every node, depth-first.

From 4d88aca8c2905aa07a7e28c41f4fe7a2c0191ea0 Mon Sep 17 00:00:00 2001
From: hongyu <linerwang1997@gmail.com>
Date: Mon, 29 Jun 2026 17:30:42 +0100
Subject: [PATCH 06/11] CHORE: typing

---
 .../scripts/nb_decomp_grouping.py             | 60 +++++++++++++++++++
 .../region_decomp/__init__.py                 | 13 +++-
 2 files changed, 71 insertions(+), 2 deletions(-)
 create mode 100644 packages/imandrax-api-models/scripts/nb_decomp_grouping.py

diff --git a/packages/imandrax-api-models/scripts/nb_decomp_grouping.py b/packages/imandrax-api-models/scripts/nb_decomp_grouping.py
new file mode 100644
index 00000000..b83bf363
--- /dev/null
+++ b/packages/imandrax-api-models/scripts/nb_decomp_grouping.py
@@ -0,0 +1,60 @@
+# %%
+from IPython.core.getipython import get_ipython
+
+if ip := get_ipython():
+    ip.run_line_magic('reload_ext', 'autoreload')
+    ip.run_line_magic('autoreload', '2')
+
+from pathlib import Path
+
+CURR_DIR = Path.cwd() if ip else Path(__file__).parent
+
+import os
+
+import dotenv
+import imandrax_api
+
+from imandrax_api_models.client import ImandraXClient
+from imandrax_api_models.region_decomp import EnrichedDecomposeRes, RegionGroup
+
+dotenv.load_dotenv()
+
+# %%
+c = ImandraXClient(
+    url=imandrax_api.url_prod,
+    # url=imandrax_api.url_dev,
+    auth_token=os.environ['IMANDRAX_API_KEY'],
+)
+
+IML = """
+let classify_triangle (a: int) (b: int) (c: int) : string =
+  (if (((a <= 0) || (b <= 0)) || (c <= 0)) then "invalid" else (if ((((a + b) <= c) || ((a + c) <= b)) || ((b + c) <= a)) then "invalid" else (if ((a = b) && (b = c)) then "equilateral" else (if (((a = b) || (b = c)) || (a = c)) then "isosceles" else "scalene"))))
+
+let is_leap_year (year: int) : bool =
+  (if ((year mod 400) = 0) then true else (if ((year mod 100) = 0) then false else ((year mod 4) = 0)))
+
+let days_in_month (year: int) (month: int) : int =
+  (if ((month < 1) || (month > 12)) then 0 else (if (month = 2) then (if (is_leap_year year) then 29 else 28) else (if ((((month = 4) || (month = 6)) || (month = 9)) || (month = 11)) then 30 else 31)))
+
+let is_valid_date (year: int) (month: int) (day: int) : bool =
+  (if (year < 1) then false else (let dim = (days_in_month year month) in
+  (if (dim = 0) then false else ((1 <= day) && (day <= dim)))))
+
+let normalize_percentage (value: real) : real =
+  (if (value <> value) then 0.0 else (if (value <. 0.0) then 0.0 else (if (value >. 100.0) then 100.0 else value)))
+[@@decomp top ~prune:true ()]
+"""
+
+_eval_res = c.eval_src(IML)
+decomp_res = c.decompose(name='normalize_percentage', prune=True, string_results=True)
+
+# %%
+print(decomp_res.regions_str)
+
+
+# %%
+hdr = EnrichedDecomposeRes.from_decomp_res(decomp_res)
+print(hdr.to_tree_str())
+
+# %%
+print(hdr.model_dump()['region_groups'])
diff --git a/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py b/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
index 335743dd..ed1c5e2a 100644
--- a/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
+++ b/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
+from collections.abc import Callable, Mapping, Sequence
 from functools import reduce
 from typing import NoReturn, Self, TypedDict
 
@@ -51,7 +51,9 @@ def _repr_html_(self) -> str:
         return mk_icicle_widget_html(self.region_groups)
 
 
-type JSONValue = str | int | float | bool | None | JSONObject | JSONArray
+type JSONValue = (
+    str | int | float | bool | None | Mapping[str, JSONValue] | Sequence[JSONValue]
+)
 type JSONObject = dict[str, JSONValue]
 type JSONArray = list[JSONValue]
 
@@ -113,6 +115,13 @@ def describe(self) -> JSONObject:
             d['example_output'] = r.model_eval_str
         return d
 
+    def to_json_dict(self) -> JSONObject:
+        """Serialize to a d3-hierarchy-compatible dict, recursing into children."""
+        d = self.describe()
+        if self.children:
+            d['children'] = [c.to_json_dict() for c in self.children]
+        return d
+
     def repr_line(self) -> str:
         """One-line representation of the region group."""
         d = self.describe()

From 4f46af5cb7ff6d9d755744f7b8f7edbd6675cfef Mon Sep 17 00:00:00 2001
From: hongyu <linerwang1997@gmail.com>
Date: Mon, 29 Jun 2026 20:26:18 +0100
Subject: [PATCH 07/11] leaf regions

---
 .../region_decomp/__init__.py                 | 50 +++++++++++++------
 .../tests/test_grouped_region_decomp.py       | 15 ++++--
 2 files changed, 44 insertions(+), 21 deletions(-)

diff --git a/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py b/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
index ed1c5e2a..90d09b04 100644
--- a/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
+++ b/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
@@ -24,7 +24,7 @@ class EnrichedDecomposeRes(DecomposeRes):
     )
 
     @model_validator(mode='after')
-    def _populate_region_groups(self) -> Self:
+    def populate_region_groups(self) -> Self:
         if not self.region_groups and self.regions_str:
             self.region_groups = group_regions(self.regions_str)
         return self
@@ -33,6 +33,16 @@ def _populate_region_groups(self) -> Self:
     def from_decomp_res(cls, v: DecomposeRes) -> EnrichedDecomposeRes:
         return cls.model_validate(v.model_dump())
 
+    @staticmethod
+    def leaf_groups(groups: list[RegionGroup]) -> list[RegionGroup]:
+        leaves = []
+        for group in groups:
+            if not group.children:
+                leaves.append(group)
+            else:
+                leaves.extend(EnrichedDecomposeRes.leaf_groups(group.children))
+        return leaves
+
     def to_tree_str(
         self,
         *,
@@ -64,27 +74,35 @@ class RegionGroup(BaseModel):
 
     Attributes:
         constraints:
-            Full accumulated constraint path from root to this node (root-first).
-            `constraints[-1]` is the constraint introduced at this node's own level.
-        label_path:
-            Positional index path from root to this node (root-first, 1-indexed).
-            Each element is the sibling index at that depth. Displayed as e.g. `1.2.3`.
-            Levels where a constraint applies to all regions are skipped, so the path
-            length may be shorter than the tree depth.
-        region:
-            The concrete region, if this group contains exactly one.
         children:
             Sub-groups under this node.
         weight:
-            Number of regions in the `has` partition at this node's level.
 
     """
 
-    constraints: list[str]
-    label_path: list[int]
-    region: RegionStr | None = None
-    children: list[RegionGroup] = []
-    weight: int
+    constraints: list[str] = Field(
+        description=(
+            'Full accumulated constraint path from root to this node (root-first).'
+            "`constraints[-1]` is the constraint introduced at this node's own level."
+        )
+    )
+    label_path: list[int] = Field(
+        description=(
+            'Positional index path from root to this node (root-first, 1-indexed).'
+            'Each element is the sibling index at that depth. Displayed as e.g. `1.2.3`.'
+            'Levels where a constraint applies to all regions are skipped, so the path'
+            'length may be shorter than the tree depth.'
+        )
+    )
+    region: RegionStr | None = Field(
+        default=None, description='The concrete region. Present iff at leaf nodes.'
+    )
+    children: list[RegionGroup] = Field(
+        default_factory=list, description='Sub-groups under this node.'
+    )
+    weight: int = Field(
+        description="Number of regions in the partition at this node's level."
+    )
 
     def n_regions(self) -> int:
         """Total regions in this subtree, including self."""
diff --git a/packages/imandrax-api-models/tests/test_grouped_region_decomp.py b/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
index 7b786167..6304c8dd 100644
--- a/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
+++ b/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
@@ -36,10 +36,15 @@ def trust() -> DecomposeRes:
 
 def test():
     decomp_res = trust()
-    hdr = EnrichedDecomposeRes.from_decomp_res(decomp_res)
+    edr = EnrichedDecomposeRes.from_decomp_res(decomp_res)
 
-    # The raw regions are inherited from DecomposeRes.
-    assert hdr.regions_str == snapshot(
+    assert edr.regions_str
+    leaf_groups = edr.leaf_groups(edr.region_groups)
+    assert len(leaf_groups) == len(edr.regions_str)
+    for leaf_group in leaf_groups:
+        assert len(leaf_group.children) == 0
+
+    assert edr.regions_str == snapshot(
         [
             RegionStr(
                 constraints_str=['y <= 0', 'x <= 0'],
@@ -81,7 +86,7 @@ def test():
     )
 
     # region_groups is auto-populated on validation from regions_str.
-    assert hdr.to_tree_str() == snapshot("""\
+    assert edr.to_tree_str() == snapshot("""\
 ├── [1] new_constraint='x >= 1' n_leaf_regions=4
 │   ├── [1.1] new_constraint='y >= 1' n_leaf_regions=2
 │   │   ├── [1.1.1] new_constraint='x <= y' invariant='2' is_leaf=True
@@ -100,7 +105,7 @@ def _walk(groups: list[RegionGroup]) -> list[tuple[str, list[str]]]:
             out.extend(_walk(g.children))
         return out
 
-    assert _walk(hdr.region_groups) == snapshot(
+    assert _walk(edr.region_groups) == snapshot(
         [
             ('1', ['x >= 1']),
             ('1.1', ['x >= 1', 'y >= 1']),

From 69370b6f69de9d0ef10939a5d93d77f10b9b7733 Mon Sep 17 00:00:00 2001
From: hongyu <linerwang1997@gmail.com>
Date: Mon, 29 Jun 2026 20:56:40 +0100
Subject: [PATCH 08/11] concrete region helper

---
 .../region_decomp/__init__.py                 | 38 +++++++++++++------
 .../tests/test_grouped_region_decomp.py       |  3 ++
 2 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py b/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
index 90d09b04..6be0cc01 100644
--- a/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
+++ b/packages/imandrax-api-models/src/imandrax_api_models/region_decomp/__init__.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Mapping, Sequence
+from dataclasses import asdict
 from functools import reduce
 from typing import NoReturn, Self, TypedDict
 
@@ -33,15 +34,18 @@ def populate_region_groups(self) -> Self:
     def from_decomp_res(cls, v: DecomposeRes) -> EnrichedDecomposeRes:
         return cls.model_validate(v.model_dump())
 
-    @staticmethod
-    def leaf_groups(groups: list[RegionGroup]) -> list[RegionGroup]:
-        leaves = []
-        for group in groups:
-            if not group.children:
-                leaves.append(group)
-            else:
-                leaves.extend(EnrichedDecomposeRes.leaf_groups(group.children))
-        return leaves
+    def regions(self) -> JSONArray:
+        """Leaf region groups (concrete regions) with hierarchical grouping info."""
+        leaf_groups = get_leaf_groups(self.region_groups)
+        ds = []
+        for leaf_group in leaf_groups:
+            d: JSONObject = {}
+            assert leaf_group.region is not None, 'Leaf group must be concrete'
+            d['label_path'] = '.'.join(map(str, leaf_group.label_path))
+            d['weight'] = leaf_group.weight
+            d |= asdict(leaf_group.region)
+            ds.append(d)
+        return ds
 
     def to_tree_str(
         self,
@@ -94,15 +98,15 @@ class RegionGroup(BaseModel):
             'length may be shorter than the tree depth.'
         )
     )
+    weight: int = Field(
+        description="Number of regions in the partition at this node's level."
+    )
     region: RegionStr | None = Field(
         default=None, description='The concrete region. Present iff at leaf nodes.'
     )
     children: list[RegionGroup] = Field(
         default_factory=list, description='Sub-groups under this node.'
     )
-    weight: int = Field(
-        description="Number of regions in the partition at this node's level."
-    )
 
     def n_regions(self) -> int:
         """Total regions in this subtree, including self."""
@@ -156,6 +160,16 @@ def repr_line(self) -> str:
         return ' '.join(parts)
 
 
+def get_leaf_groups(groups: list[RegionGroup]) -> list[RegionGroup]:
+    leaves = []
+    for group in groups:
+        if not group.children:
+            leaves.append(group)
+        else:
+            leaves.extend(get_leaf_groups(group.children))
+    return leaves
+
+
 def group_regions(regions: list[RegionStr]) -> list[RegionGroup]:
     """Group regions hierarchically based on constraints."""
     return _loop_group_regions([], [], regions)
diff --git a/packages/imandrax-api-models/tests/test_grouped_region_decomp.py b/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
index 6304c8dd..e3e45aec 100644
--- a/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
+++ b/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
@@ -43,6 +43,9 @@ def test():
     assert len(leaf_groups) == len(edr.regions_str)
     for leaf_group in leaf_groups:
         assert len(leaf_group.children) == 0
+        assert leaf_group.region, 'Leaf group must be concrete'
+        assert leaf_group.region.constraints_str
+        assert set(leaf_group.constraints) == set(leaf_group.region.constraints_str)
 
     assert edr.regions_str == snapshot(
         [

From dc188b643bb90a2a8043cbd58bbab103865c2421 Mon Sep 17 00:00:00 2001
From: hongyu <linerwang1997@gmail.com>
Date: Mon, 29 Jun 2026 20:56:49 +0100
Subject: [PATCH 09/11] fmt enriched regions

---
 .../src/imandrax_api_models/context_utils.py    | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/packages/imandrax-api-models/src/imandrax_api_models/context_utils.py b/packages/imandrax-api-models/src/imandrax_api_models/context_utils.py
index c8955789..1c55fe60 100644
--- a/packages/imandrax-api-models/src/imandrax_api_models/context_utils.py
+++ b/packages/imandrax-api-models/src/imandrax_api_models/context_utils.py
@@ -15,6 +15,7 @@
     Position,
     VerifyRes,
 )
+from imandrax_api_models.region_decomp import EnrichedDecomposeRes
 from imandrax_api_models.yaml_utils import ImandraXAPIModelDumper
 
 
@@ -317,3 +318,19 @@ def format_decomp_res(decomp_res: DecomposeRes) -> str:
 
     data = remove_art_and_task_fields(data)
     return yaml.dump(data, Dumper=ImandraXAPIModelDumper, width=120)
+
+
+def format_enriched_decomp_res(decomp_res: EnrichedDecomposeRes) -> dict[str, Any]:
+    # TOOD: use remove_fields_rec
+    d: dict[str, Any] = {}
+    if decomp_res.regions_str is not None:
+        enriched_regions = decomp_res.regions()
+        d['descr'] = f'Decomp succeeded with {len(enriched_regions)} regions'
+        d['regions'] = enriched_regions
+    else:
+        d['descr'] = 'Decomp failed'
+        d |= remove_art_and_task_fields(decomp_res.model_dump())
+        d.pop('regions_str')
+        d.pop('region_groups')
+
+    return d

From 64d9155df3864264f6ca53bba5eddbcca48641af Mon Sep 17 00:00:00 2001
From: hongyu <linerwang1997@gmail.com>
Date: Mon, 29 Jun 2026 20:58:54 +0100
Subject: [PATCH 10/11] TEST,FIX: imports

---
 .../tests/test_grouped_region_decomp.py                   | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/packages/imandrax-api-models/tests/test_grouped_region_decomp.py b/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
index e3e45aec..f59ebadf 100644
--- a/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
+++ b/packages/imandrax-api-models/tests/test_grouped_region_decomp.py
@@ -5,7 +5,11 @@
 from inline_snapshot import snapshot
 
 from imandrax_api_models.proto_models import DecomposeRes
-from imandrax_api_models.region_decomp import EnrichedDecomposeRes, RegionGroup
+from imandrax_api_models.region_decomp import (
+    EnrichedDecomposeRes,
+    RegionGroup,
+    get_leaf_groups,
+)
 
 
 def trust() -> DecomposeRes:
@@ -39,7 +43,7 @@ def test():
     edr = EnrichedDecomposeRes.from_decomp_res(decomp_res)
 
     assert edr.regions_str
-    leaf_groups = edr.leaf_groups(edr.region_groups)
+    leaf_groups = get_leaf_groups(edr.region_groups)
     assert len(leaf_groups) == len(edr.regions_str)
     for leaf_group in leaf_groups:
         assert len(leaf_group.children) == 0

From e018bd2703bdef5f58082750a7b2a6023b15262a Mon Sep 17 00:00:00 2001
From: hongyu <linerwang1997@gmail.com>
Date: Mon, 29 Jun 2026 21:01:21 +0100
Subject: [PATCH 11/11] CLEANUP

---
 .../scripts/nb_decomp_grouping.py             | 60 -------------------
 1 file changed, 60 deletions(-)
 delete mode 100644 packages/imandrax-api-models/scripts/nb_decomp_grouping.py

diff --git a/packages/imandrax-api-models/scripts/nb_decomp_grouping.py b/packages/imandrax-api-models/scripts/nb_decomp_grouping.py
deleted file mode 100644
index b83bf363..00000000
--- a/packages/imandrax-api-models/scripts/nb_decomp_grouping.py
+++ /dev/null
@@ -1,60 +0,0 @@
-# %%
-from IPython.core.getipython import get_ipython
-
-if ip := get_ipython():
-    ip.run_line_magic('reload_ext', 'autoreload')
-    ip.run_line_magic('autoreload', '2')
-
-from pathlib import Path
-
-CURR_DIR = Path.cwd() if ip else Path(__file__).parent
-
-import os
-
-import dotenv
-import imandrax_api
-
-from imandrax_api_models.client import ImandraXClient
-from imandrax_api_models.region_decomp import EnrichedDecomposeRes, RegionGroup
-
-dotenv.load_dotenv()
-
-# %%
-c = ImandraXClient(
-    url=imandrax_api.url_prod,
-    # url=imandrax_api.url_dev,
-    auth_token=os.environ['IMANDRAX_API_KEY'],
-)
-
-IML = """
-let classify_triangle (a: int) (b: int) (c: int) : string =
-  (if (((a <= 0) || (b <= 0)) || (c <= 0)) then "invalid" else (if ((((a + b) <= c) || ((a + c) <= b)) || ((b + c) <= a)) then "invalid" else (if ((a = b) && (b = c)) then "equilateral" else (if (((a = b) || (b = c)) || (a = c)) then "isosceles" else "scalene"))))
-
-let is_leap_year (year: int) : bool =
-  (if ((year mod 400) = 0) then true else (if ((year mod 100) = 0) then false else ((year mod 4) = 0)))
-
-let days_in_month (year: int) (month: int) : int =
-  (if ((month < 1) || (month > 12)) then 0 else (if (month = 2) then (if (is_leap_year year) then 29 else 28) else (if ((((month = 4) || (month = 6)) || (month = 9)) || (month = 11)) then 30 else 31)))
-
-let is_valid_date (year: int) (month: int) (day: int) : bool =
-  (if (year < 1) then false else (let dim = (days_in_month year month) in
-  (if (dim = 0) then false else ((1 <= day) && (day <= dim)))))
-
-let normalize_percentage (value: real) : real =
-  (if (value <> value) then 0.0 else (if (value <. 0.0) then 0.0 else (if (value >. 100.0) then 100.0 else value)))
-[@@decomp top ~prune:true ()]
-"""
-
-_eval_res = c.eval_src(IML)
-decomp_res = c.decompose(name='normalize_percentage', prune=True, string_results=True)
-
-# %%
-print(decomp_res.regions_str)
-
-
-# %%
-hdr = EnrichedDecomposeRes.from_decomp_res(decomp_res)
-print(hdr.to_tree_str())
-
-# %%
-print(hdr.model_dump()['region_groups'])