diff --git a/src/mccode_antlr/instr/instr.py b/src/mccode_antlr/instr/instr.py index 74a2815..6111977 100644 --- a/src/mccode_antlr/instr/instr.py +++ b/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]) diff --git a/tests/instr/test_split.py b/tests/instr/test_split.py index 7b763c4..a8ea3df 100644 --- a/tests/instr/test_split.py +++ b/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'})