algo2code 0.2.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (46) hide show
  1. algo2code-0.2.1/.gitignore +75 -0
  2. algo2code-0.2.1/LICENSE +21 -0
  3. algo2code-0.2.1/PKG-INFO +30 -0
  4. algo2code-0.2.1/README.md +13 -0
  5. algo2code-0.2.1/prototypes/__init__.py +54 -0
  6. algo2code-0.2.1/prototypes/algo_parser.py +471 -0
  7. algo2code-0.2.1/prototypes/ast_nodes.py +153 -0
  8. algo2code-0.2.1/prototypes/demo.py +98 -0
  9. algo2code-0.2.1/prototypes/expr_parser.py +491 -0
  10. algo2code-0.2.1/prototypes/main.py +90 -0
  11. algo2code-0.2.1/prototypes/pcg.tex +34 -0
  12. algo2code-0.2.1/prototypes/pcg_generated.py +78 -0
  13. algo2code-0.2.1/prototypes/taichi_codegen.py +528 -0
  14. algo2code-0.2.1/prototypes/test_algo2code.py +390 -0
  15. algo2code-0.2.1/prototypes/type_inference.py +209 -0
  16. algo2code-0.2.1/pyproject.toml +36 -0
  17. algo2code-0.2.1/src/algo2code/__init__.py +72 -0
  18. algo2code-0.2.1/src/algo2code/algo_parser.py +458 -0
  19. algo2code-0.2.1/src/algo2code/ast_nodes.py +171 -0
  20. algo2code-0.2.1/src/algo2code/backends/__init__.py +1 -0
  21. algo2code-0.2.1/src/algo2code/backends/c_petsc_codegen.py +1 -0
  22. algo2code-0.2.1/src/algo2code/backends/numpy_codegen.py +1 -0
  23. algo2code-0.2.1/src/algo2code/backends/taichi_codegen.py +859 -0
  24. algo2code-0.2.1/src/algo2code/errors.py +27 -0
  25. algo2code-0.2.1/src/algo2code/expr_parser.py +598 -0
  26. algo2code-0.2.1/src/algo2code/library/.gitkeep +0 -0
  27. algo2code-0.2.1/src/algo2code/library/__init__.py +19 -0
  28. algo2code-0.2.1/src/algo2code/library/pcg.py +127 -0
  29. algo2code-0.2.1/src/algo2code/library/radial_return_j2.py +90 -0
  30. algo2code-0.2.1/src/algo2code/library/radial_return_j2_kinematic.py +89 -0
  31. algo2code-0.2.1/src/algo2code/library/radial_return_j2_mixed.py +91 -0
  32. algo2code-0.2.1/src/algo2code/type_inference.py +219 -0
  33. algo2code-0.2.1/tests/conftest.py +45 -0
  34. algo2code-0.2.1/tests/test_algo_parser.py +60 -0
  35. algo2code-0.2.1/tests/test_deferrals.py +65 -0
  36. algo2code-0.2.1/tests/test_end_to_end.py +22 -0
  37. algo2code-0.2.1/tests/test_expr_parser.py +134 -0
  38. algo2code-0.2.1/tests/test_fail_loud.py +92 -0
  39. algo2code-0.2.1/tests/test_radial_return_codegen.py +116 -0
  40. algo2code-0.2.1/tests/test_radial_return_j2_kinematic_codegen.py +136 -0
  41. algo2code-0.2.1/tests/test_radial_return_j2_mixed_codegen.py +175 -0
  42. algo2code-0.2.1/tests/test_smoke.py +43 -0
  43. algo2code-0.2.1/tests/test_taichi_codegen.py +64 -0
  44. algo2code-0.2.1/tests/test_transpose_alias.py +42 -0
  45. algo2code-0.2.1/tests/test_type_inference.py +32 -0
  46. algo2code-0.2.1/tests/test_vector_lowering.py +148 -0
@@ -0,0 +1,75 @@
1
+ # Python
2
+ __pycache__/
3
+ *.py[cod]
4
+ *$py.class
5
+ *.egg-info/
6
+ *.egg
7
+ dist/
8
+ build/
9
+ *.whl
10
+ .eggs/
11
+
12
+ # Virtual environments
13
+ .venv/
14
+ venv/
15
+ env/
16
+
17
+ # IDE
18
+ .idea/
19
+ .vscode/
20
+ *.swp
21
+ *.swo
22
+ *~
23
+
24
+ # Testing
25
+ .pytest_cache/
26
+ .coverage
27
+ htmlcov/
28
+ .mypy_cache/
29
+ .ruff_cache/
30
+
31
+ # OS
32
+ .DS_Store
33
+ Thumbs.db
34
+
35
+ # Taichi
36
+ *.tcb
37
+ ti_cache/
38
+
39
+ # Artifacts
40
+ *.artifact.json
41
+ *.artifact.yaml
42
+
43
+ # Jupyter
44
+ .ipynb_checkpoints/
45
+
46
+ # Claude Code
47
+ .claude/settings.local.json
48
+
49
+ # Claude CoWork
50
+ /MechDsl/
51
+ *.npz
52
+ # Golden-file regression snapshots must be tracked
53
+ !packages/mechdsl-core/tests/golden/*.npz
54
+ .gitnexus
55
+
56
+ # Local working artifacts (orchestra runtime, plan-edit backups)
57
+ .orchestra/
58
+ **/.orchestra/
59
+ *.original.md
60
+ *.original.[0-9].md
61
+ .claude/worktrees/
62
+ .publisher-worktrees/
63
+ .sesskey
64
+
65
+ # Public-release staging tree (curated copy pushed to CEmM2/MechDSL)
66
+ /dev/MechDSL/
67
+
68
+ # MkDocs build output
69
+ /site/
70
+
71
+ # Generated wiki / code-intelligence index (logic-loom / akms tooling)
72
+ .repo_wiki/
73
+
74
+ # comment-sweep working tree
75
+ /.comment-review/
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Shmuel Osovski
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,30 @@
1
+ Metadata-Version: 2.4
2
+ Name: algo2code
3
+ Version: 0.2.1
4
+ Summary: Transpile LaTeX algorithm boxes (algpseudocode) to executable Taichi/NumPy/C code
5
+ Project-URL: Repository, https://github.com/CEmM2/MechDSL
6
+ Author: Shmuel Osovski
7
+ License-Expression: MIT
8
+ License-File: LICENSE
9
+ Keywords: algorithm,code-generation,latex,taichi,transpiler
10
+ Classifier: Development Status :: 3 - Alpha
11
+ Classifier: Intended Audience :: Science/Research
12
+ Classifier: License :: OSI Approved :: MIT License
13
+ Classifier: Programming Language :: Python :: 3
14
+ Classifier: Programming Language :: Python :: 3.12
15
+ Requires-Python: <3.14,>=3.11
16
+ Description-Content-Type: text/markdown
17
+
18
+ # algo2code
19
+
20
+ Transpile LaTeX algorithm boxes (`algpseudocode`) to executable code targeting Taichi, NumPy, or C/PETSc.
21
+
22
+ Parses `\begin{algorithmic}...\end{algorithmic}` environments with type-directed code generation.
23
+ Zero runtime dependencies — standard library only.
24
+
25
+ ## Documentation
26
+
27
+ - **User docs:** <https://sosovski.group/MechDSL/algo2code/> — introduction, getting
28
+ started, usage, and examples.
29
+ - **Design spec:** `dev/design_docs/11-ALGO2CODE.md` (authoritative source of truth).
30
+ - See the [monorepo root](../../README.md) for the full project overview.
@@ -0,0 +1,13 @@
1
+ # algo2code
2
+
3
+ Transpile LaTeX algorithm boxes (`algpseudocode`) to executable code targeting Taichi, NumPy, or C/PETSc.
4
+
5
+ Parses `\begin{algorithmic}...\end{algorithmic}` environments with type-directed code generation.
6
+ Zero runtime dependencies — standard library only.
7
+
8
+ ## Documentation
9
+
10
+ - **User docs:** <https://sosovski.group/MechDSL/algo2code/> — introduction, getting
11
+ started, usage, and examples.
12
+ - **Design spec:** `dev/design_docs/11-ALGO2CODE.md` (authoritative source of truth).
13
+ - See the [monorepo root](../../README.md) for the full project overview.
@@ -0,0 +1,54 @@
1
+ """
2
+ algo2code — LaTeX algorithmic environment → executable code transpiler.
3
+
4
+ Usage:
5
+ from algo2code import transpile
6
+
7
+ taichi_code = transpile(latex_source, backend='taichi')
8
+ """
9
+ from .ast_nodes import Algorithm, VarType
10
+ from .algo_parser import parse_algorithm
11
+ from .expr_parser import parse_latex_expr
12
+ from .type_inference import infer_types
13
+ from .taichi_codegen import generate_taichi
14
+
15
+
16
+ def transpile(source: str, backend: str = 'taichi') -> str:
17
+ """
18
+ Full pipeline: LaTeX source → parsed AST → type inference → code generation.
19
+
20
+ Parameters
21
+ ----------
22
+ source : str
23
+ LaTeX source containing \\begin{algorithmic} ... \\end{algorithmic}
24
+ with optional % directive comments.
25
+ backend : str
26
+ Target backend. Currently only 'taichi' is supported.
27
+
28
+ Returns
29
+ -------
30
+ str
31
+ Generated source code.
32
+ """
33
+ # 1. Parse
34
+ algo = parse_algorithm(source)
35
+
36
+ # 2. Type inference
37
+ infer_types(algo)
38
+
39
+ # 3. Code generation
40
+ if backend == 'taichi':
41
+ return generate_taichi(algo)
42
+ else:
43
+ raise ValueError(f"Unknown backend: {backend!r}. Supported: 'taichi'")
44
+
45
+
46
+ __all__ = [
47
+ 'transpile',
48
+ 'parse_algorithm',
49
+ 'parse_latex_expr',
50
+ 'infer_types',
51
+ 'generate_taichi',
52
+ 'Algorithm',
53
+ 'VarType',
54
+ ]
@@ -0,0 +1,471 @@
1
+ r"""
2
+ Parser for LaTeX `algpseudocode` environments.
3
+
4
+ Recognises:
5
+ \begin{algorithmic} ... \end{algorithmic}
6
+ \State $lhs = rhs$ % optional_comment
7
+ \For{$k = 0, 1, \ldots, N$} ... \EndFor
8
+ \While{$cond$} ... \EndWhile
9
+ \If{$cond$} ... \ElsIf{$cond$} ... \Else ... \EndIf
10
+ \Return $expr$
11
+ \State \textbf{break}
12
+
13
+ Directive comments (lines starting with `%` OUTSIDE algorithmic):
14
+ % algorithm <name>
15
+ % backend <taichi|numpy|petsc>
16
+ % args <name>:<type>, ...
17
+ % type <varname> <scalar|vector|matrix|callable>
18
+
19
+ The parser delegates all $...$ math fragments to expr_parser.
20
+ """
21
+ from __future__ import annotations
22
+ import re
23
+ from .ast_nodes import (
24
+ Algorithm, Stmt, Assign, ForLoop, WhileLoop, Branch, Return, Break,
25
+ Var, VarType
26
+ )
27
+ from .expr_parser import parse_latex_expr, parse_assignment, parse_condition
28
+
29
+
30
+ # ── Directive parsing ────────────────────────────────────────────────────────
31
+
32
+ _TYPE_MAP = {
33
+ 'scalar': VarType.SCALAR,
34
+ 'vector': VarType.VECTOR,
35
+ 'matrix': VarType.MATRIX,
36
+ 'callable': VarType.CALLABLE,
37
+ 'matvec': VarType.MATRIX, # alias
38
+ }
39
+
40
+
41
+ def _parse_directives(lines: list[str]) -> dict:
42
+ """Parse % directive comments before the algorithmic block."""
43
+ directives: dict = {
44
+ 'name': 'algorithm',
45
+ 'backend': 'taichi',
46
+ 'args': [],
47
+ 'types': {},
48
+ }
49
+
50
+ for line in lines:
51
+ line = line.strip()
52
+ if not line.startswith('%'):
53
+ continue
54
+ line = line[1:].strip()
55
+
56
+ if line.startswith('algorithm '):
57
+ directives['name'] = line.split(None, 1)[1].strip()
58
+
59
+ elif line.startswith('backend '):
60
+ directives['backend'] = line.split(None, 1)[1].strip()
61
+
62
+ elif line.startswith('args '):
63
+ arg_str = line.split(None, 1)[1]
64
+ for arg in arg_str.split(','):
65
+ arg = arg.strip()
66
+ if ':' in arg:
67
+ name, typ = arg.split(':', 1)
68
+ directives['args'].append(
69
+ (name.strip(), _TYPE_MAP.get(typ.strip(), VarType.UNKNOWN))
70
+ )
71
+ else:
72
+ directives['args'].append((arg, VarType.UNKNOWN))
73
+
74
+ elif line.startswith('type '):
75
+ parts = line.split()
76
+ if len(parts) >= 3:
77
+ varname = parts[1]
78
+ vtype = _TYPE_MAP.get(parts[2], VarType.UNKNOWN)
79
+ directives['types'][varname] = vtype
80
+
81
+ return directives
82
+
83
+
84
+ # ── Main parser ──────────────────────────────────────────────────────────────
85
+
86
+ class AlgPseudocodeParser:
87
+ """
88
+ Parse a complete LaTeX source containing an algorithmic environment.
89
+
90
+ Usage:
91
+ algo = AlgPseudocodeParser(latex_string).parse()
92
+ """
93
+
94
+ def __init__(self, source: str):
95
+ self.source = source
96
+ self.lines: list[str] = []
97
+ self.pos = 0
98
+
99
+ def parse(self) -> Algorithm:
100
+ """Parse the full source and return an Algorithm AST."""
101
+ # Split into pre-algorithmic directives and body
102
+ pre_lines, body_lines = self._split_sections()
103
+ directives = _parse_directives(pre_lines)
104
+
105
+ # Also collect inline type comments from body lines
106
+ type_annotations = dict(directives['types'])
107
+ self._collect_inline_types(body_lines, type_annotations)
108
+
109
+ # Parse the body
110
+ self.lines = body_lines
111
+ self.pos = 0
112
+ stmts = self._parse_block(terminators=[])
113
+
114
+ return Algorithm(
115
+ name=directives['name'],
116
+ backend=directives['backend'],
117
+ args=directives['args'],
118
+ body=stmts,
119
+ type_annotations=type_annotations,
120
+ )
121
+
122
+ def _split_sections(self) -> tuple[list[str], list[str]]:
123
+ """Split source into pre-algorithmic lines and body lines."""
124
+ all_lines = self.source.split('\n')
125
+ pre_lines = []
126
+ body_lines = []
127
+ in_body = False
128
+
129
+ for line in all_lines:
130
+ stripped = line.strip()
131
+
132
+ if re.match(r'\\begin\{algorithmic\}', stripped):
133
+ in_body = True
134
+ continue
135
+ if re.match(r'\\end\{algorithmic\}', stripped):
136
+ in_body = False
137
+ continue
138
+ # Also handle \begin{algorithm} wrapper
139
+ if re.match(r'\\begin\{algorithm\}', stripped):
140
+ continue
141
+ if re.match(r'\\end\{algorithm\}', stripped):
142
+ continue
143
+ if re.match(r'\\caption\{', stripped):
144
+ continue
145
+
146
+ if in_body:
147
+ if stripped: # skip blank lines
148
+ body_lines.append(line)
149
+ else:
150
+ pre_lines.append(line)
151
+
152
+ return pre_lines, body_lines
153
+
154
+ def _collect_inline_types(self, lines: list[str], types: dict):
155
+ """Extract inline type annotations from % comments on \State lines."""
156
+ for line in lines:
157
+ # Look for: \State $...$ % vector
158
+ # or: \State $...$ % type_for:varname=vector
159
+ m = re.search(r'%\s*(\w+)\s*$', line)
160
+ if m:
161
+ hint = m.group(1).lower()
162
+ if hint in _TYPE_MAP:
163
+ # Infer which variable this annotates from the LHS
164
+ lhs_match = re.search(r'\\State\s+\$\s*([^=\$]+?)\s*=', line)
165
+ if lhs_match:
166
+ var_name = self._extract_var_name(lhs_match.group(1))
167
+ if var_name:
168
+ types[var_name] = _TYPE_MAP[hint]
169
+
170
+ def _extract_var_name(self, lhs_latex: str) -> str | None:
171
+ """Extract a clean variable name from a LaTeX LHS fragment."""
172
+ # Strip \mathbf{...}, \boldsymbol{...}, etc.
173
+ lhs = re.sub(r'\\(?:mathbf|boldsymbol|bm)\{([^}]*)\}', r'\1', lhs_latex)
174
+ # Strip \\text{...} subscripts for naming
175
+ lhs = re.sub(r'_\{\\text\{([^}]*)\}\}', r'_\1', lhs)
176
+ lhs = re.sub(r'\\(\w+)', r'\1', lhs) # \alpha -> alpha
177
+ lhs = lhs.strip().replace(' ', '')
178
+ return lhs if lhs else None
179
+
180
+ # ── Statement parsing ────────────────────────────────────────────────
181
+
182
+ def _current_line(self) -> str | None:
183
+ if self.pos < len(self.lines):
184
+ return self.lines[self.pos].strip()
185
+ return None
186
+
187
+ def _advance(self):
188
+ self.pos += 1
189
+
190
+ def _parse_block(self, terminators: list[str]) -> list[Stmt]:
191
+ """Parse statements until a terminator command is found."""
192
+ stmts = []
193
+ while self.pos < len(self.lines):
194
+ line = self._current_line()
195
+ if line is None:
196
+ break
197
+
198
+ # Check if this line starts with any terminator
199
+ stripped = self._strip_comment(line)
200
+ if any(stripped.startswith(t) for t in terminators):
201
+ break
202
+
203
+ stmt = self._parse_statement()
204
+ if stmt is not None:
205
+ stmts.append(stmt)
206
+
207
+ return stmts
208
+
209
+ def _strip_comment(self, line: str) -> str:
210
+ """Remove trailing % comment but preserve % inside $...$."""
211
+ in_math = False
212
+ for i, ch in enumerate(line):
213
+ if ch == '$':
214
+ in_math = not in_math
215
+ elif ch == '%' and not in_math:
216
+ return line[:i].strip()
217
+ return line.strip()
218
+
219
+ def _extract_inline_comment(self, line: str) -> str:
220
+ """Extract the % comment portion."""
221
+ in_math = False
222
+ for i, ch in enumerate(line):
223
+ if ch == '$':
224
+ in_math = not in_math
225
+ elif ch == '%' and not in_math:
226
+ return line[i + 1:].strip()
227
+ return ''
228
+
229
+ def _parse_statement(self) -> Stmt | None:
230
+ """Parse a single statement from the current line."""
231
+ line = self._current_line()
232
+ if line is None:
233
+ return None
234
+
235
+ stripped = self._strip_comment(line)
236
+ comment = self._extract_inline_comment(line)
237
+
238
+ # ── \For{...} ──
239
+ if stripped.startswith('\\For'):
240
+ return self._parse_for(stripped)
241
+
242
+ # ── \While{...} ──
243
+ if stripped.startswith('\\While'):
244
+ return self._parse_while(stripped)
245
+
246
+ # ── \If{...} ──
247
+ if stripped.startswith('\\If'):
248
+ return self._parse_if(stripped)
249
+
250
+ # ── \Return ──
251
+ if stripped.startswith('\\Return') or stripped.startswith('\\State \\Return'):
252
+ self._advance()
253
+ return self._parse_return(stripped)
254
+
255
+ # ── \State \textbf{break} or \State \Break ──
256
+ if re.search(r'\\textbf\{break\}|\\Break|\\textbf\{Break\}', stripped):
257
+ self._advance()
258
+ return Break()
259
+
260
+ # ── \State $assignment$ ──
261
+ if stripped.startswith('\\State'):
262
+ self._advance()
263
+ return self._parse_state(stripped, comment)
264
+
265
+ # Skip unrecognized lines
266
+ self._advance()
267
+ return None
268
+
269
+ def _extract_math(self, text: str) -> str:
270
+ """Extract content between $ delimiters."""
271
+ m = re.search(r'\$(.+?)\$', text)
272
+ if m:
273
+ return m.group(1).strip()
274
+ # Try without $ (some formats omit them)
275
+ return text.strip()
276
+
277
+ def _extract_brace_arg(self, text: str, command: str) -> str:
278
+ """Extract the {argument} after a \\Command."""
279
+ # Find the command, then extract balanced braces
280
+ idx = text.find(command)
281
+ if idx < 0:
282
+ return ''
283
+ rest = text[idx + len(command):]
284
+
285
+ # Find opening brace
286
+ brace_start = rest.find('{')
287
+ if brace_start < 0:
288
+ return ''
289
+
290
+ depth = 0
291
+ start = brace_start
292
+ for i in range(brace_start, len(rest)):
293
+ if rest[i] == '{':
294
+ depth += 1
295
+ elif rest[i] == '}':
296
+ depth -= 1
297
+ if depth == 0:
298
+ return rest[start + 1:i].strip()
299
+ return rest[start + 1:].strip()
300
+
301
+ # ── For loop ─────────────────────────────────────────────────────────
302
+
303
+ def _parse_for(self, line: str) -> ForLoop:
304
+ """Parse \\For{$k = 0, 1, \\ldots, N$} body \\EndFor"""
305
+ arg = self._extract_brace_arg(line, '\\For')
306
+ arg = arg.strip('$ ')
307
+
308
+ var, start, end_expr = self._parse_for_range(arg)
309
+
310
+ self._advance() # past the \For line
311
+ body = self._parse_block(terminators=['\\EndFor'])
312
+
313
+ # Consume the \EndFor line
314
+ if self._current_line() and self._strip_comment(self._current_line()).startswith('\\EndFor'):
315
+ self._advance()
316
+
317
+ return ForLoop(var=var, start=start, end_expr=end_expr, body=body)
318
+
319
+ def _parse_for_range(self, arg: str) -> tuple[str, int, str]:
320
+ """
321
+ Parse for-loop range specifications:
322
+ k = 0, 1, ..., N → var='k', start=0, end='N'
323
+ k = 0, 1, 2, ... → var='k', start=0, end=''
324
+ k = 1 to N → var='k', start=1, end='N'
325
+ """
326
+ # Pattern: var = start, ..., end
327
+ m = re.match(
328
+ r'([a-zA-Z]\w*)\s*=\s*(\d+)\s*,\s*\d+\s*,?\s*'
329
+ r'(?:\\ldots|\\dots|\\cdots|\.\.\.)\s*(?:,\s*)?'
330
+ r'(?:\\(?:text|mathrm)\{(\w+)\}|([a-zA-Z]\w*))?',
331
+ arg
332
+ )
333
+ if m:
334
+ var = m.group(1)
335
+ start = int(m.group(2))
336
+ end_expr = m.group(3) or m.group(4) or ''
337
+ return var, start, end_expr
338
+
339
+ # Pattern: var = start to end
340
+ m = re.match(r'([a-zA-Z]\w*)\s*=\s*(\d+)\s+(?:to|\\to)\s+(\w+)', arg)
341
+ if m:
342
+ return m.group(1), int(m.group(2)), m.group(3)
343
+
344
+ # Fallback
345
+ m = re.match(r'([a-zA-Z]\w*)', arg)
346
+ var = m.group(1) if m else 'k'
347
+ return var, 0, ''
348
+
349
+ # ── While loop ───────────────────────────────────────────────────────
350
+
351
+ def _parse_while(self, line: str) -> WhileLoop:
352
+ arg = self._extract_brace_arg(line, '\\While')
353
+ arg = arg.strip('$ ')
354
+ condition = parse_condition(arg)
355
+
356
+ self._advance()
357
+ body = self._parse_block(terminators=['\\EndWhile'])
358
+
359
+ if self._current_line() and self._strip_comment(self._current_line()).startswith('\\EndWhile'):
360
+ self._advance()
361
+
362
+ return WhileLoop(condition=condition, body=body)
363
+
364
+ # ── If / ElsIf / Else ────────────────────────────────────────────────
365
+
366
+ def _parse_if(self, line: str) -> Branch:
367
+ arg = self._extract_brace_arg(line, '\\If')
368
+ arg = arg.strip('$ ')
369
+ condition = parse_condition(arg)
370
+
371
+ self._advance()
372
+ if_body = self._parse_block(
373
+ terminators=['\\EndIf', '\\ElsIf', '\\Else']
374
+ )
375
+
376
+ elif_branches = []
377
+ else_body = []
378
+
379
+ while self._current_line():
380
+ cur = self._strip_comment(self._current_line())
381
+ if cur.startswith('\\ElsIf'):
382
+ elif_arg = self._extract_brace_arg(cur, '\\ElsIf')
383
+ elif_arg = elif_arg.strip('$ ')
384
+ elif_cond = parse_condition(elif_arg)
385
+ self._advance()
386
+ elif_body = self._parse_block(
387
+ terminators=['\\EndIf', '\\ElsIf', '\\Else']
388
+ )
389
+ elif_branches.append((elif_cond, elif_body))
390
+ elif cur.startswith('\\Else'):
391
+ self._advance()
392
+ else_body = self._parse_block(terminators=['\\EndIf'])
393
+ break
394
+ else:
395
+ break
396
+
397
+ if self._current_line() and self._strip_comment(self._current_line()).startswith('\\EndIf'):
398
+ self._advance()
399
+
400
+ return Branch(
401
+ condition=condition,
402
+ if_body=if_body,
403
+ elif_branches=elif_branches,
404
+ else_body=else_body,
405
+ )
406
+
407
+ # ── Return ───────────────────────────────────────────────────────────
408
+
409
+ def _parse_return(self, line: str) -> Return:
410
+ # Extract everything after \Return
411
+ m = re.search(r'\\Return\s*(.*)', self._strip_comment(line))
412
+ if not m:
413
+ return Return(values=[])
414
+
415
+ rest = m.group(1).strip().strip('$').strip()
416
+ if not rest:
417
+ return Return(values=[])
418
+
419
+ # Parse comma-separated return values
420
+ values = []
421
+ for part in self._split_top_level(rest, ','):
422
+ part = part.strip()
423
+ if part:
424
+ values.append(parse_latex_expr(part))
425
+
426
+ return Return(values=values)
427
+
428
+ def _split_top_level(self, text: str, sep: str) -> list[str]:
429
+ """Split text by separator, respecting brace depth."""
430
+ parts = []
431
+ depth = 0
432
+ current = []
433
+ for ch in text:
434
+ if ch == '{':
435
+ depth += 1
436
+ elif ch == '}':
437
+ depth -= 1
438
+ elif ch == sep and depth == 0:
439
+ parts.append(''.join(current))
440
+ current = []
441
+ continue
442
+ current.append(ch)
443
+ parts.append(''.join(current))
444
+ return parts
445
+
446
+ # ── State (assignment) ───────────────────────────────────────────────
447
+
448
+ def _parse_state(self, line: str, comment: str) -> Stmt | None:
449
+ """Parse \\State $lhs = rhs$"""
450
+ stripped = self._strip_comment(line)
451
+ stripped = re.sub(r'^\\State\s*', '', stripped).strip()
452
+ # Extract math content
453
+ math = self._extract_math(stripped)
454
+ if not math:
455
+ return None
456
+
457
+ result = parse_assignment(math)
458
+ if result is None:
459
+ # Not an assignment — could be a standalone expression
460
+ expr = parse_latex_expr(math)
461
+ return Assign(target=Var(name='_'), value=expr, comment=comment)
462
+
463
+ target, value = result
464
+ return Assign(target=target, value=value, comment=comment)
465
+
466
+
467
+ # ── Public API ───────────────────────────────────────────────────────────────
468
+
469
+ def parse_algorithm(source: str) -> Algorithm:
470
+ """Parse a LaTeX source containing an algorithmic environment."""
471
+ return AlgPseudocodeParser(source).parse()