mcpl_split drops USERVARS blocks a part does not use

src/mccode_antlr/instr/instr.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 from functools import lru_cache
+import re
 from io import StringIO
 from os import PathLike
 from pathlib import Path
@@ -946,6 +947,10 @@ def mcpl_split(self,
         # move the newly added component to the front of the list:
         second.components = (second.components[-1],) + second.components[:-1]
 
+        if remove_unused_parameters:
+            # Only now: the MCPL components' parameters can name particle variables too
+            first.remove_unused_user_vars()
+            second.remove_unused_user_vars()
         return first, second
 
     def parameter_used(self, name: str):
@@ -961,6 +966,41 @@ def parameter_used(self, name: str):
                     return True
         return False
 
+    def remove_unused_user_vars(self) -> int:
+        """Remove the USERVARS blocks declaring only names nothing else here mentions.
+
+        A split instrument keeps the particle variables of the part it lost, so the same
+        primary spectrometer differs with every sample that declares its own. A name counts
+        as mentioned if it appears anywhere in the instrument or its components' code, as
+        a word: in code, or quoted in a parameter such as ``tof="event_time"``. A block
+        whose declarations can not be read is kept.
+
+        Returns
+        -------
+        int
+            The number of USERVARS blocks removed
+        """
+        if not self.user:
+            return 0
+        user, self.user = self.user, ()
+        text = StringIO()
+        try:
+            self.to_file(text)
+        finally:
+            self.user = user
+        words = set(_WORD.findall(text.getvalue()))
+        for comp in {id(x.type): x.type for x in self.components}.values():
+            for name in _COMP_BLOCKS:
+                for block in getattr(comp, name):
+                    words.update(_WORD.findall(str(block)))
+        kept = tuple(block for block in self.user
+                     if not (names := _declared_names(str(block))) or not names.isdisjoint(words))
+        removed = len(self.user) - len(kept)
+        if removed:
+            logger.info(f'Removed {removed} USERVARS block(s) of {self.name} it does not use')
+            self.user = kept
+        return removed
+
     def check_instrument_parameters(self, remove=False):
         """Check that all instrument parameters are used in the instrument, and optionally remove any that are not
 
@@ -995,6 +1035,21 @@ def check_expr(self, expr: int | float | str | Expr):
         return expr
 
 
+_WORD = re.compile(r'[A-Za-z_]\w*')
+_COMP_BLOCKS = ('share', 'user', 'declare', 'initialize', 'trace', 'save', 'final', 'display')
+
+
+def _declared_names(source: str) -> set[str]:
+    """The names a USERVARS block declares: ``int a, b[3]; double *c;`` -> a, b, c"""
+    source = re.sub(r'//[^\n]*|/\*.*?\*/', '', source, flags=re.S)
+    names = set()
+    for declaration in source.split(';'):
+        for declarator in re.sub(r'\[[^]]*]', '', declaration).split(','):
+            if found := re.search(r'([A-Za-z_]\w*)\s*$', declarator):
+                names.add(found.group(1))
+    return names
+
+
 def _join_raw_tuple(raw_tuple: tuple[RawC, ...]):
     return '\n'.join([str(rc) for rc in raw_tuple])
 

tests/instr/test_split.py

@@ -117,3 +117,46 @@ def test_split_broken_reference(self):
             self.assertEqual(name, s_inst.name)
             self.assertTrue(s_inst.at_relative[1] is not None)
             self.assertTrue(s_inst.rotate_relative[1] is not None)
+
+
+class TestSplitUserVars(TestCase):
+    @staticmethod
+    def instr():
+        """Each user_vars call adds its own block, as a sample or readout section does."""
+        from mccode_antlr.assembler import Assembler
+        a = Assembler('test_uservars')
+        a.user_vars('double incident_speed;')
+        a.user_vars('int sample_switch; // which process scatters\ndouble weights[3], *spare;')
+        a.user_vars('int RING;')
+        a.component('origin', 'Arm', at=(0, 0, 0)).EXTEND('incident_speed = sqrt(vx*vx + vy*vy + vz*vz);')
+        a.component('split_here', 'Arm', at=(0, 0, 1))
+        a.component('sample', 'Arm', at=(0, 0, 2)).EXTEND('sample_switch = rand01() > 0.5;')
+        a.component('detector', 'Arm', at=(0, 0, 3)).WHEN('RING > 0')
+        return a.instrument
+
+    def blocks(self, instr):
+        return [' '.join(str(b).split()) for b in instr.user]
+
+    def test_split_removes_user_vars_a_part_does_not_use(self):
+        first, second = self.instr().mcpl_split('split_here', remove_unused_parameters=True)
+        self.assertEqual(self.blocks(first), ['double incident_speed;'])
+        self.assertEqual(self.blocks(second), [
+            'int sample_switch; // which process scatters double weights[3], *spare;',
+            'int RING;'])
+
+    def test_split_keeps_user_vars_the_mcpl_components_name(self):
+        from mccode_antlr.common import ComponentParameter, Expr
+        flag = (ComponentParameter('userflag', Expr.string('"RING"')),)
+        first, _ = self.instr().mcpl_split('split_here', output_parameters=flag,
+                                           remove_unused_parameters=True)
+        self.assertEqual(self.blocks(first), ['double incident_speed;', 'int RING;'])
+
+    def test_split_keeps_user_vars_unless_asked(self):
+        first, second = self.instr().mcpl_split('split_here')
+        self.assertEqual(len(first.user), 3)
+        self.assertEqual(len(second.user), 3)
+
+    def test_declared_names(self):
+        from mccode_antlr.instr.instr import _declared_names
+        self.assertEqual(_declared_names('int a, b[3]; /* c; */ double *d; // e;\nchar f[2][4];'),
+                         {'a', 'b', 'd', 'f'})