From 93d1cddd270ac797c74eeca7a4b259d2edfabb99 Mon Sep 17 00:00:00 2001
From: Claude <claude@gitea.local>
Date: Fri, 16 Oct 2026 03:27:13 +0000
Subject: [PATCH] BlockLayout2D.build: check capacity before generating regular
 origins

c397e08 moved the 64-bit capacity check after np.arange(nby*nbx), so an
over-limit regular layout allocated before rejecting. Check arithmetic
block count (regular) or validated selected count (sparse) first; mock
regression asserts no np.arange before the ValueError.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/CelerisLab/lbm/blocks.py    | 16 ++++++++++------
 tests/unit/test_block_layout.py |  7 +++++++
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/src/CelerisLab/lbm/blocks.py b/src/CelerisLab/lbm/blocks.py
index 8869fe2..5fe109f 100644
--- a/src/CelerisLab/lbm/blocks.py
+++ b/src/CelerisLab/lbm/blocks.py
@@ -140,13 +140,17 @@ class BlockLayout2D:
         nby, nbx = ny // by, nx // bx
         ly, lx = by + 2 * halo, bx + 2 * halo
         cells = ly * lx
-        if origins is None:
-            jy, jx = np.divmod(np.arange(nby * nbx), nbx)
-            origins = np.stack([jy * by, jx * bx], axis=1)
-        selected = _selected_blocks(origins, ny, nx, by, bx)
-        origins = np.array(origins, dtype=np.int64).reshape(-1, 2)  # own copy, never caller storage
-        if (len(origins) + 1) * NQ * cells > np.iinfo(np.int64).max // 2:
+        if origins is not None:
+            selected = _selected_blocks(origins, ny, nx, by, bx)
+            origins = np.array(origins, dtype=np.int64).reshape(-1, 2)  # own copy, never caller storage
+        # Capacity from arithmetic before any grid-sized allocation; hole =
+        # nblocks * cells is below this bound too.
+        if (nby * nbx if origins is None else len(origins)) * NQ * cells > np.iinfo(np.int64).max // 2:
             raise ValueError("blocked storage exceeds 64-bit indexing")
+        if origins is None:
+            selected = np.arange(nby * nbx)
+            jy, jx = np.divmod(selected, nbx)
+            origins = np.stack([jy * by, jx * bx], axis=1)
 
         # Grid block -> stored block, or len(origins) for an absent block, so
         # absent cells land exactly on hole = nblocks * cells.
diff --git a/tests/unit/test_block_layout.py b/tests/unit/test_block_layout.py
index 282931c..a3a686b 100644
--- a/tests/unit/test_block_layout.py
+++ b/tests/unit/test_block_layout.py
@@ -1,6 +1,7 @@
 """CPU block layout contracts: ownership, halo pairs, packing and tagged pull."""
 import sys
 import unittest
+from unittest import mock
 from types import SimpleNamespace
 
 import numpy as np
@@ -280,6 +281,12 @@ class BlockLayoutTests(unittest.TestCase):
         lx = layout.local_shape[1]
         self.assertEqual(int(mapped.write[0]), int(layout.slot(3 * lx + 0, 1)))
 
+    def test_capacity_rejected_before_allocation(self):
+        fail = AssertionError("grid-sized allocation before the capacity check")
+        with mock.patch.object(np, "arange", side_effect=fail):
+            with self.assertRaises(ValueError):
+                BlockLayout2D.build(2 ** 30, 2 ** 30, 1, 1)
+
     def test_tables_are_readonly(self):
         layout = BlockLayout2D.build(6, 8, 3, 4)
         for array in (layout.origins, layout.owner, layout.halo_dst, layout.halo_src, layout.halo_kind):