Fix tao.init design lattice file.bmad@use_line-style syntax

src/latform/dump.py

@@ -222,7 +222,10 @@ def _resolve_used_elements(
     """
     Resolve which named items are active in the expanded lattice.
 
-    Starting from the roots of the last USE statement, lines are expanded
+    Roots are resolved per top-level lattice file: the last USE statement in
+    each file's call tree, or the ``@use_line`` names from the corresponding
+    ``tao.init`` ``design_lattice(i)%file`` entry, which override that
+    lattice's own USE statement.  From those roots, lines are expanded
     recursively (including repetitions, reflections, replacement-line calls,
     list members, and fork targets).  A fixpoint pass then folds in usage that
     depends on other elements being used:
@@ -239,17 +242,17 @@ def _resolve_used_elements(
         Mapping of uppercased active names to a human-readable reason.
     """
     all_statements = files.get_statements_in_order(repeat_called_files=False)
-    use_cmds = [
-        st
-        for st in all_statements
-        if isinstance(st, Simple) and st.statement._upper == "USE" and st.arguments
-    ]
 
-    # Bmad semantics: the last USE statement wins; each of its arguments is
-    # the root line of a branch.  Bare argument tokens are parsed as
-    # value-less Attributes.
-    roots: list[Token] = []
-    if use_cmds:
+    def last_use_roots(statements: list[Statement]) -> list[Token]:
+        # Bmad semantics: the last USE statement wins; each of its arguments is
+        # the root line of a branch.  Bare argument tokens are parsed as
+        # value-less Attributes.
+        use_cmds = [
+            st
+            for st in statements
+            if isinstance(st, Simple) and st.statement._upper == "USE" and st.arguments
+        ]
+        roots: list[Token] = []
         for use_cmd in reversed(use_cmds):
             for arg in use_cmd.arguments:
                 if isinstance(arg, Token):
@@ -264,6 +267,27 @@ def _resolve_used_elements(
             # I think this is a scenario that's common in reused sublattices
             # that can be standalone
             break
+        return roots
+
+    # Roots are resolved per top-level lattice file.  A tao.init
+    # ``design_lattice(i)%file = 'lat.bmad@line_name'`` suffix overrides that
+    # lattice's own USE statement (Tao/bmad_parser semantics); lattices without
+    # a suffix fall back to their in-file USE.
+    tao_entries: list[tuple[str, list[str]]] = (
+        files.tao_init.lattice_file_with_use_line if files.tao_init is not None else []
+    )
+    root_reasons: dict[Token, str] = {}
+    for index, top_file in enumerate(files.top_files):
+        use_lines = tao_entries[index][1] if index < len(tao_entries) else []
+        if use_lines:
+            for name in use_lines:
+                root_reasons.setdefault(Token(name.upper()), "tao.init use_line")
+        else:
+            statements = files.get_statements_in_order(
+                repeat_called_files=False, top_files=[top_file]
+            )
+            for root in last_use_roots(statements):
+                root_reasons.setdefault(root, "use statement")
 
     deferred = get_deferred_element_attributes(all_statements)
     used: dict[Token, str] = {}
@@ -288,10 +312,10 @@ def mark(name: Token, reason: str) -> None:
                 for target in _fork_targets(item, deferred):
                     mark(target, f"fork target of {name}")
 
-    for root in roots:
-        mark(root, "use statement")
+    for root, reason in root_reasons.items():
+        mark(root, reason)
 
-    if roots:
+    if root_reasons:
         # The implicit lattice endpoints exist in any expanded lattice
         for name in (Token("BEGINNING"), Token("END")):
             if name in named_items:
@@ -319,7 +343,7 @@ def mark(name: Token, reason: str) -> None:
                     changed = True
                 continue
 
-            if not roots:
+            if not root_reasons:
                 continue
 
             superposition = element.get_superposition_settings(deferred)

src/latform/parser.py

@@ -944,8 +944,18 @@ def flatten_all(self, call: bool, inline: bool) -> dict[pathlib.Path, list[State
         """Flatten each top-level file independently, keyed by its path."""
         return {top: self.flatten(call=call, inline=inline, top=top) for top in self.top_files}
 
-    def get_statements_in_order(self, *, repeat_called_files: bool = True) -> list[Statement]:
-        """Get all statements in order, as evaluated by Bmad."""
+    def get_statements_in_order(
+        self,
+        *,
+        repeat_called_files: bool = True,
+        top_files: Sequence[pathlib.Path] | None = None,
+    ) -> list[Statement]:
+        """
+        Get all statements in order, as evaluated by Bmad.
+
+        ``top_files`` restricts the traversal to the call trees of the given
+        top-level files (defaulting to all of them).
+        """
 
         handled = set()
         active = set()
@@ -969,7 +979,7 @@ def _flatten(fn):
                 active.discard(fn)
 
         res = []
-        for fn in self.top_files:
+        for fn in self.top_files if top_files is None else top_files:
             res.extend(_flatten(fn))
         return res
 

src/latform/tao/file.py

@@ -421,21 +421,45 @@ def _lattice_file_assignments(self) -> list[Assignment]:
         ]
 
     @property
-    def lattice_files(self) -> list[str]:
+    def lattice_file_with_use_line(self) -> list[tuple[str, list[str]]]:
         """
-        Ordered ``design_lattice(i)%file`` values (unquoted), by index.
+        Ordered, unquoted ``design_lattice(i)%file`` values with `use_line` names.
 
-        When set, rewritse the ``design_lattice(i)%file`` entries to ``files`` (1-based).
-        Existing entries are updated in place. Non-matching additional entries
-        are removed, and new entries are appended.
+        When set, rewrites the entries from ``(filename, use_line_names)``
+        tuples, joining the names back into the ``@`` suffix form.
         """
+
+        def item(assignment: Assignment) -> tuple[str, list[str]]:
+            # Unquote and remove the '@use_line1@use_line2...' suffix
+
+            filename = unquote_value(assignment.value.strip())
+            if "@" in filename:
+                filename, use_lines = filename.split("@", 1)
+                return (filename, use_lines.split("@"))
+            return filename, []
+
         by_index = {
-            assignment.path.components[0].index: unquote_value(assignment.value.strip())
+            assignment.path.components[0].index: item(assignment)
             for assignment in self._lattice_file_assignments()
             if assignment.path.components[0].index is not None
         }
         return [by_index[i] for i in sorted(by_index)]
 
+    @lattice_file_with_use_line.setter
+    def lattice_file_with_use_line(self, entries: list[tuple[str, list[str]]]) -> None:
+        self.lattice_files = ["@".join([fn, *use_lines]) for fn, use_lines in entries]
+
+    @property
+    def lattice_files(self) -> list[str]:
+        """
+        Ordered ``design_lattice(i)%file`` values (unquoted), by index.
+
+        When set, rewrites the ``design_lattice(i)%file`` entries to ``files`` (1-based).
+        Existing entries are updated in place. Non-matching additional entries
+        are removed, and new entries are appended.
+        """
+        return [fn for fn, _ in self.lattice_file_with_use_line]
+
     @lattice_files.setter
     def lattice_files(self, files: list[str]) -> None:
         namelist = self.design_lattice

src/latform/template.py

@@ -253,17 +253,6 @@ def _instantiate_tao_init(
 
     in_dir = posixpath.dirname(input_rel)
     out_dir = posixpath.dirname(output_rel)
-    remapped: list[str] = []
-    changed = False
-    for entry in tao_init.lattice_files:
-        mapped = _map_reference(entry, in_dir, out_dir, explicit_paths, in_to_out)
-        if mapped is None:
-            remapped.append(entry)
-        else:
-            remapped.append(mapped)
-            changed = True
-    if changed:
-        tao_init.lattice_files = remapped
 
     merged: dict[str, str] = {}
     if rules and (rules["literal"] or rules["regex"] or rules["parts"]):
@@ -276,6 +265,20 @@ def _instantiate_tao_init(
     if merged:
         rename_tao_elements(tao_init, merged)
 
+    remapped: list[tuple[str, list[str]]] = []
+    changed = False
+    for filename, use_lines in tao_init.lattice_file_with_use_line:
+        mapped = _map_reference(filename, in_dir, out_dir, explicit_paths, in_to_out)
+        mapped = mapped or filename
+
+        new_use_lines = [merged.get(name.upper(), name) for name in use_lines]
+        remapped.append((mapped, new_use_lines))
+
+        if mapped != filename or new_use_lines != use_lines:
+            changed = True
+    if changed:
+        tao_init.lattice_file_with_use_line = remapped
+
     for name_key, assignments in ((override or {}).get("namelists") or {}).items():
         name, index = split_namelist_key(name_key)
         interpolated = {k: _interpolate(str(v), instance) for k, v in assignments.items()}

src/latform/tests/test_dump.py

@@ -227,6 +227,69 @@ def test_reasons():
     assert status["UNUSED_Q"]["reason"] == ""
 
 
+TAO_TWO_LINE_LATTICE = """\
+m1: marker
+m2: marker
+l1: line = (m1)
+l2: line = (m2)
+use, l1
+"""
+
+
+def _tao_files(entries: list[str], lattice_contents: dict[str, str]) -> MemoryFiles:
+    body = "\n".join(
+        f"  design_lattice({i})%file = '{entry}'" for i, entry in enumerate(entries, start=1)
+    )
+    files = MemoryFiles.from_tao_init_contents(
+        f"&tao_design_lattice\n{body}\n/\n",
+        "/virtual/tao.init",
+        lattice_contents=lattice_contents,
+    )
+    files.parse()
+    files.annotate()
+    return files
+
+
+def _tao_used(entries: list[str], lattice_contents: dict[str, str]) -> set[str]:
+    files = _tao_files(entries, lattice_contents)
+    return {row["name"] for row in get_elements_status(files, "all") if row["used"] == "YES"}
+
+
+@pytest.mark.parametrize(
+    ("entry", "expected_used", "expected_unused"),
+    [
+        pytest.param("mem.lat.bmad", {"L1", "M1"}, {"L2", "M2"}, id="no-use-line"),
+        pytest.param("mem.lat.bmad@l2", {"L2", "M2"}, {"L1", "M1"}, id="use-line-overrides-use"),
+        pytest.param(
+            "mem.lat.bmad@l1@l2", {"L1", "M1", "L2", "M2"}, set(), id="multiple-use-lines"
+        ),
+    ],
+)
+def test_tao_init_use_line_roots(entry: str, expected_used: set[str], expected_unused: set[str]):
+    used = _tao_used([entry], {"mem.lat.bmad": TAO_TWO_LINE_LATTICE})
+    assert expected_used <= used
+    assert not expected_unused & used
+
+
+def test_tao_init_use_line_is_per_file():
+    """A use_line overrides only its own lattice's USE; other lattices keep theirs."""
+    first = "m1: marker\nm2: marker\nl1: line = (m1)\nl2: line = (m2)\nuse, l1\n"
+    second = "m3: marker\nm4: marker\nl3: line = (m3)\nl4: line = (m4)\nuse, l3\n"
+    used = _tao_used(
+        ["first.lat.bmad", "second.lat.bmad@l4"],
+        {"first.lat.bmad": first, "second.lat.bmad": second},
+    )
+    assert {"L1", "M1", "L4", "M4"} <= used
+    assert not {"L2", "M2", "L3", "M3"} & used
+
+
+def test_tao_init_use_line_reason():
+    files = _tao_files(["mem.lat.bmad@l2"], {"mem.lat.bmad": TAO_TWO_LINE_LATTICE})
+    status = {row["name"]: row for row in get_elements_status(files, "all")}
+    assert status["L2"]["reason"] == "tao.init use_line"
+    assert status["M2"]["reason"] == "in line L2"
+
+
 def test_get_constants():
     src = """
 my_const = 0.5

src/latform/tests/test_parse.py

@@ -2745,6 +2745,32 @@ def test_roundtrip(code: str, expected) -> None:
     roundtrip_code(code)
 
 
+@pytest.mark.parametrize(
+    "code",
+    [
+        pytest.param("q: quad, l = ((1+2))*3", id="doubled-parens"),
+        pytest.param("ov: overlay = {{q[k1]: 0.1}}, var = {a}", id="doubled-braces"),
+        pytest.param(
+            "foo[phi0_multipass] = -(acos(e2/e1*cos(phi/180*pi)) + phi/180*pi)/(2*pi)",
+            id="negated-group",
+        ),
+        pytest.param(
+            "bar[phi0_multipass] = -(\n  acos(e2/e1*cos(phi/180*pi)) + phi/180*pi\n)/(2*pi)",
+            id="negated-group-multiline",
+        ),
+    ],
+)
+def test_nested_parens_roundtrip(code: str) -> None:
+    # Redundant nesting like (( )) is collapsed on output, so this checks
+    # parse -> format -> parse equality rather than byte preservation.
+    roundtrip_code(code)
+
+
+def test_nested_brackets_invalid() -> None:
+    with pytest.raises(ValueError, match=r"Nested '\[\[ \]\]'"):
+        parse("m: marker, type = [[x]]")
+
+
 @pytest.mark.parametrize(
     ("code",),
     [

src/latform/tests/test_tao_init.py

@@ -88,6 +88,44 @@ def test_lattice_files_ordered_and_skips_comments():
     ]
 
 
+USE_LINE_INIT = """\
+&tao_design_lattice
+  design_lattice(2)%file = 'b.lat.bmad@lineA'
+  design_lattice(1)%file = 'a.lat.bmad'
+  design_lattice(3)%file = 'sub/c.lat.bmad@lineA@lineB'
+/
+"""
+
+
+def test_lattice_file_with_use_line():
+    tao = TaoInit.parse(USE_LINE_INIT)
+    assert tao.lattice_file_with_use_line == [
+        ("a.lat.bmad", []),
+        ("b.lat.bmad", ["lineA"]),
+        ("sub/c.lat.bmad", ["lineA", "lineB"]),
+    ]
+    # The '@use_line' suffixes are stripped from the plain filename view
+    assert tao.lattice_files == ["a.lat.bmad", "b.lat.bmad", "sub/c.lat.bmad"]
+    assert tao.render() == USE_LINE_INIT
+
+
+def test_set_lattice_file_with_use_line():
+    tao = TaoInit.parse(USE_LINE_INIT)
+    entries = [("a2.lat.bmad", []), ("b2.lat.bmad", ["lineC"]), ("c2.lat.bmad", ["l1", "l2"])]
+    tao.lattice_file_with_use_line = entries
+    assert tao.lattice_file_with_use_line == entries
+    assert tao.lattice_files == ["a2.lat.bmad", "b2.lat.bmad", "c2.lat.bmad"]
+    assert TaoInit.parse(tao.render()).lattice_file_with_use_line == entries
+
+
+def test_set_lattice_files_passes_at_suffix_through():
+    """The plain setter writes '@' suffixes verbatim; the getter strips them (asymmetric)."""
+    tao = TaoInit.parse(USE_LINE_INIT)
+    tao.lattice_files = ["x.lat.bmad@ln"]
+    assert tao.lattice_files == ["x.lat.bmad"]
+    assert tao.lattice_file_with_use_line == [("x.lat.bmad", ["ln"])]
+
+
 def test_keypath_decomposes_nested_key():
     path = KeyPath.parse("foo(3)%bar(2)%val")
     assert path.names == ("foo", "bar", "val")
@@ -250,6 +288,21 @@ def test_memory_files_from_tao_init_contents():
     assert "M_Q" in files.get_named_items()
 
 
+def test_memory_files_from_tao_init_use_line():
+    """A '@use_line' suffix does not interfere with resolving the lattice file itself."""
+    contents = "&tao_design_lattice\n  design_lattice(1)%file = 'mem.lat.bmad@ml'\n/\n"
+    root = FILES / "virtual" / "tao.init"
+    files = MemoryFiles.from_tao_init_contents(
+        contents,
+        root,
+        lattice_contents={"mem.lat.bmad": "M_Q: quadrupole, l = 1\nml: line = (M_Q)\nuse, ml\n"},
+    )
+    assert [p.name for p in files.top_files] == ["mem.lat.bmad"]
+    files.parse()
+    files.annotate()
+    assert "M_Q" in files.get_named_items()
+
+
 @pytest.mark.parametrize(
     ("name", "expected"),
     [("tao.init", True), ("tao_plot.init", True), ("foo.INIT", True), ("lat.bmad", False)],