fix(data): enforce semantic chunk claims

Author: apstenku123

tests/test_dedup_store_chunk_claims.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import sqlite3
+import sys
+from pathlib import Path
+
+
+MLX_ROOT = Path(__file__).resolve().parents[1]
+CLANG_INDEXER = MLX_ROOT / "tools" / "clang_indexer"
+if str(CLANG_INDEXER) not in sys.path:
+    sys.path.insert(0, str(CLANG_INDEXER))
+
+
+def test_chunk_claims_are_wal_backed_and_count_limited(tmp_path):
+    from dedup_store import DedupStore
+
+    db = tmp_path / "dedup.sqlite"
+    first = DedupStore(str(db), near=False, commit_every=1)
+    second = DedupStore(str(db), near=False, commit_every=1)
+    try:
+        assert first.conn.execute("PRAGMA journal_mode").fetchone()[0] == "wal"
+
+        assert first.claim_chunk_tokens([10, 20, 30], namespace="strict") is True
+        assert second.claim_chunk_tokens([10, 20, 30], namespace="strict") is False
+
+        assert first.claim_chunk_tokens([10, 20, 30], namespace="repeat3", max_count=3)
+        assert second.claim_chunk_tokens([10, 20, 30], namespace="repeat3", max_count=3)
+        assert first.claim_chunk_tokens([10, 20, 30], namespace="repeat3", max_count=3)
+        assert not second.claim_chunk_tokens(
+            [10, 20, 30],
+            namespace="repeat3",
+            max_count=3,
+        )
+    finally:
+        first.close()
+        second.close()
+
+    conn = sqlite3.connect(db)
+    try:
+        rows = dict(
+            conn.execute(
+                "SELECT namespace, claim_count FROM chunk_claims ORDER BY namespace"
+            ).fetchall()
+        )
+    finally:
+        conn.close()
+
+    assert rows == {"repeat3": 3, "strict": 1}

tools/clang_indexer/dedup_store.py

@@ -185,6 +185,13 @@ def _init_schema(self) -> None:
         self._execute_write(
             "CREATE TABLE IF NOT EXISTS dedup_meta (key TEXT PRIMARY KEY, val INTEGER)"
         )
+        self._execute_write(
+            "CREATE TABLE IF NOT EXISTS chunk_claims ("
+            "namespace TEXT NOT NULL, "
+            "hash BLOB NOT NULL, "
+            "claim_count INTEGER NOT NULL, "
+            "PRIMARY KEY(namespace, hash))"
+        )
         # next doc_id counter for near-dup docs.
         self._execute_write(
             "INSERT OR IGNORE INTO dedup_meta (key, val) VALUES ('next_doc_id', 0)"
@@ -369,6 +376,44 @@ def seen_near_tokens(self, token_ids: Sequence[int]) -> bool:
         self._mark_pending()
         return False
 
+    def claim_chunk_tokens(
+        self,
+        token_ids: Sequence[int],
+        *,
+        namespace: str = "train_chunk",
+        max_count: int = 1,
+    ) -> bool:
+        """Claim one semantic training chunk by its tokenized body.
+
+        This is separate from the ``exact``/``minhash`` dedup tables. Exact/near
+        dedup decides whether a function may be emitted as a root document;
+        chunk claims decide whether that already-tokenized function/class/type
+        body may appear anywhere in the training stream. The unit of ownership is
+        still the caller-provided semantic chunk; the hash is only the SQLite key.
+
+        Returns True when this call successfully claimed one slot. Returns False
+        when the tokenized chunk has already reached ``max_count`` in ``namespace``.
+        The write is committed immediately so concurrent conveyor processes see
+        the claim before they assemble overlapping 1k/2k/4k/8k buckets.
+        """
+        if max_count < 1:
+            raise ValueError("claim_chunk_tokens requires max_count >= 1")
+        if not namespace:
+            raise ValueError("claim_chunk_tokens requires a non-empty namespace")
+        h = _sha1_tokens(token_ids)
+        cur = self._execute_write(
+            "INSERT INTO chunk_claims(namespace, hash, claim_count) "
+            "VALUES (?, ?, 1) "
+            "ON CONFLICT(namespace, hash) DO UPDATE SET "
+            "claim_count = claim_count + 1 "
+            "WHERE claim_count < ?",
+            (namespace, h, int(max_count)),
+        )
+        # Claims are coordination records, not buffered dedup rows. Commit them
+        # immediately so other WAL connections see the slot before emitting docs.
+        self.commit()
+        return cur.rowcount == 1
+
     # ----------------------------------------------------------------- #
     def _load_minhash(self, doc_id: int):
         row = self.conn.execute(