Beta implementation of Just In time parsing by MicahGale · Pull Request #960 · idaholab/MontePy

MicahGale · 2026-05-29T13:58:48Z

Pull Request Checklist for MontePy

Description

This PR adds the ability to do Just-in-time parsing. The quick summary is that at file read time, only the very beginning (3 words max per object) are read and used to store the object in the "proper place" then when more information is needed from that object it is full_parsed, and the data are stored so it may be accessed. Here's an overview of how that was done:

Complete rework of __init__ to a hook based system with _init_blank, _parse_tree, etc. To make the code more DRY.
Added jit_parse args to __init__ and defaulted to jit_parse.
Added full_parse that essentially triggers a re-init
Added the decorators needs_full_cst and needs_full_ast to mark functions that need full abstract/concrete syntax trees. Right now they are identical, but they could be used in the future. When a function that is marked is called it is checked if it is fully parsed, if not a full parse occurs.
Added decorators to pull objects from other places in the problem on the fly to remove the need for update_pointers

For more details read the updated dev docs.

I know this is a ton @tjlaboss. I am thinking of having you review just the architecture, and then calling that good before doing a beta release. For testing I found a way to have all of the integration tests to run with both jit_parse=False and jit_parse=True pretty cleanly with pytest (thanks Claude).

TODO:

Get 100% patch coverage.

Fixes #529

General Checklist

I have performed a self-review of my own code.
The code follows the standards outlined in the development documentation.
I have formatted my code with black version 25 or 26.
I have added tests that prove my fix is effective or that my feature works (if applicable).

LLM Disclosure

Are you?
- A human user
- A large language model (LLM), including ones acting on behalf of a human
Were any large language models (LLM or "AI") used in to generate any of this code?

Yes
- Model(s) used: Claude (Sonnet 4.6), GH Copilot (?)
No

Documentation Checklist

I have documented all added classes and methods.
For infrastructure updates, I have updated the developer's guide.
For significant new features, I have added a section to the getting started guide.

First-Time Contributor Checklist

If this is your first contribution, add yourself to pyproject.toml if you wish to do so.

Additional Notes for Reviewers

Ensure that:

This PR fully addresses and resolves the referenced issue(s).
The submitted code is consistent with the merge checklist outlined here.
The PR covers all relevant aspects according to the development guidelines.
100% coverage of the patch is achieved, or justification for a variance is given.

📚 Documentation preview 📚: https://montepy--960.org.readthedocs.build/en/960/

tjlaboss · 2026-06-24T15:59:22Z

Observation:

JIT significantly reduces memory usage--less memory allocation and more garbage collection

So far JIT somewhat inflates loading time. Without JIT: 31 s. With JIT: 69s, for the benchmark. Reading a large reactor model with depletion on my desktop, loading time increased from 41 to 46 s.

MicahGale · 2026-07-02T19:59:15Z

So far JIT somewhat inflates loading time. Without JIT: 31 s. With JIT: 69s, for the benchmark. Reading a large reactor model with depletion on my desktop, loading time increased from 41 to 46 s.

I suspect the overhead is due to @args_checked I will dig into this more.

re.match only checks at the start of a string, so patterns would miss keywords anywhere past the first character of a line. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Previously all JIT cells were unconditionally full-parsed to check for redundant cell-block definitions. Now Cell.search() is used as a cheap pre-filter: only cells whose raw text contains the modifier keyword are full-parsed, and only then is set_in_cell_block consulted. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

In JIT mode finalize_init now calls _check_redundant_definitions instead of push_to_cells for data-block cell modifiers, keeping those modifiers lazy until their data is actually needed. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

…ate. When a parked value is applied during full_parse the negative-sign flag (which encodes not_truncated) was silently dropped because _accept_and_update only set _old_number. Now it also sets _not_truncated = value.is_negative. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Three related fixes for correct JIT universe handling: 1. universe_input._tree_value: resolve self.universe before setting val.value to avoid corrupting _old_number. The old order set val.value = 0 first, which caused old_number to return 0, making @prop_pointer_from_problem look up Universe(0) instead of the real universe. That false lookup triggered grab_cells_from_jit_parse for Universe(0), which search-matched unrelated cells and wrongly claimed them (e.g., cell 99 being stolen from Universe(350)). 2. universe_input.push_to_cells: correctly detect ValueNode(None) (a J jump entry) as a jump via an explicit is_jump predicate, and check hasattr(cell, "_not_parsed") instead of hasattr(cell._universe, "_not_parsed") in the jump/no-data else branch. Blank cell modifiers always have _not_parsed=True even for non-JIT reads, so the old check was incorrectly skipping non-JIT cells with jump entries, preventing Universe(0) from being created during push_to_cells. 3. universes.Universes.__getitem__: persist the universe returned by _find_and_populate_universe into the collection. Without this, lazily-found universes had _collection=None, so number-conflict checks on universe.number = X would silently no-op. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

…ut a number token Previously every JIT-parsed data input received classifier.number = ValueNode(None, int), making the classifier appear to carry a number even when none was present in the input. Only assign classifier.number when a NUMBER or NULL token was actually consumed. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

…e directly _surface_type is populated by the JIT light parser, so there is no need to go through the surface_type property (which is @needs_full_ast). Also honour the jit_parse=False contract for unrecognised surface types by calling full_parse() before returning the buffer surface. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

…aterial assembly Using the old_number property or the thermal_scattering setter on a JIT object forces a full parse at read time. Access _old_number.value and _thermal_scattering directly in Materials._append_hook. Also override ThermalScatteringLaw.format_for_mcnp_input to force a full parse before formatting so the output is correct regardless of JIT state. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

…fs and guard sub-collection ownership After deepcopy, _collection_ref weakrefs on member objects point to the original (pre-copy) collection because weakrefs are not deep-copied. __setstate__ now re-establishes them against the new collection instance. Also tighten __internal_append so it only claims ownership of an object when there is no existing problem-level collection owner. This prevents a cell's _surfaces or _complements sub-collection from displacing the authoritative problem.surfaces/problem.cells owner during append. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

…n_data_block, keep _part_combos across full parse Three related fixes in CellModifierInput: 1. Leading-comment preservation: _jit_light_init inserts a start_pad PaddingNode into the JIT tree so _grab_beginning_comment fires and captures comments that appear before the data-block input. _original_lines prepends those nodes when returning raw lines, and full_parse restores them into the freshly-initialised tree so they survive reinit. 2. print_in_data_block awareness: format_for_mcnp_input now checks print_in_data_block in the JIT fast-path and falls through to a full parse (triggering push_to_cells) when the modifier belongs in the opposite block from where it is being printed. 3. _part_combos preservation: the _KEYS_TO_PRESERVE guard switches from an "is not None" check to a truthy check so an empty _part_combos list from the JIT state does not overwrite the freshly-parsed [{n,p}] value after full_parse reinitialises the object. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

…push_to_cells already registered it Accessing self.data (decorated @needs_full_ast) inside _find_and_populate_universe triggers full_parse() -> push_to_cells(), which creates and registers the target universe in problem.universes. The method then created a second Universe with the same number, causing a NumberConflictError when Universes.__getitem__ tried to append it. Check universes.get(number) after the found-detection block and return the existing object if push_to_cells already added it. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

For a Union member that's a parameterised generic (e.g. list[Particle] in `list[Particle] | set[Particle]`), check_type couldn't isinstance() against the whole union at once, so it tried each member in turn via try/except TypeError. Every call passing a set through that annotation (e.g. Importance._generate_default_data_tree, once per particle per cell) paid for a full recursive per-element validation against the wrong branch, a raised-and-caught TypeError, before succeeding on the right one -- 1002 wasted raises parsing benchmark/big_model.imcnp. Add _union_shape_matches to cheaply test a member's container type (via typing.get_origin, unwrapping Annotated first) before doing the expensive per-element check, so only the actually-matching branch gets fully validated. Falls back to the try/except loop only when multiple candidates share the same container type (e.g. list[int] | list[str]), which still needs disambiguation. Confirmed via instrumentation this eliminates all 1002 wasted TypeErrors on the benchmark model, with the full test suite passing unchanged. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

…apping Every call through an @args_checked-wrapped function paid for inspect.Signature.bind(), which re-derives the argument-name-to-value mapping from scratch on every invocation using a general state machine that also re-validates call arity, duplicate bindings, etc. -- all things the real call to func() below will validate identically. This was the single largest source of overhead in args_checked: 93,180 calls parsing benchmark/big_model.imcnp spent a combined ~0.4s of self-time inside Signature.bind()/_bind() alone. Precompute the signature shape once at decoration time (positional parameter order, and the names of any *args/**kwargs parameters), then at call time map each positional/keyword argument straight to its checkers with a plain loop. A malformed call (wrong arity, unknown keyword, etc.) simply isn't checked here and falls through to raise its own TypeError from the real call, exactly as bind() would have raised, just without duplicating that validation up front -- and with one fewer frame in the resulting traceback. Confirmed via cProfile that Signature.bind/_bind no longer appear in the hot path at all, with the full test suite and doctests passing unchanged. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Every attribute write on an MCNP_Object subclass previously did a hasattr() followed by a getattr() on the class to check for a property descriptor. Both walk the MRO, so this was two lookups where one suffices. Collapse to a single getattr(..., None) call; behavior is unchanged, including for underscore-prefixed properties with setters (e.g. Importance._explicitly_set). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Cherry-picked/ported from f6cdc6d on meta_optimize (develop), where this metaclass lives inline in mcnp_object.py; here it's its own module, so porting it required a manual edit instead of a clean cherry-pick. Same dead-code bug: the `if key.startswith("_")` branch never short-circuited the loop, so the unconditional `inspect.isfunction(value)` check right after it immediately overwrote the decision, wrapping every method and property -- public, private, and dunder -- in a try/except purely to attach file/line context to exceptions. Add the missing `continue`, keeping `__init__` as the one wrapped exception since MCNP_Problem.parse_input isn't wrapped by this metaclass either. On this JIT-heavy branch the effect is much larger than on develop: wrapped calls parsing benchmark/big_model.imcnp drop from 590,086 to 30,066 (-95%), since JIT parsing leans far more heavily on private helper methods (_jit_light_init, _parse_tree, etc.) than the classic eager-parse path does. Full test suite passes unchanged. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

_SurfaceClassFactory generates a per-subclass __init__ (used by e.g. XPlane, Sphere) that does nothing but immediately forward input, number, and jit_parse unchanged into Surface.__init__, which has identical annotations for all three and performs the real validation right after. Confirmed via runtime instrumentation this was one of the "double guarded" checks found while profiling: ~1002 wasted check_type_and_value calls parsing benchmark/big_model.imcnp, one pair per surface. Surface.__init__ stays checked -- it's the shared choke-point every surface subclass routes through, including hand-written subclasses whose spec has more than one surface_type and so never get this auto-generated wrapper at all. Only the redundant outer hop is removed; a malformed call still raises the same TypeError one frame later, from Surface.__init__ itself. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

parse_data's own @args_checked was redundant: its very first line unconditionally constructs data_input.DataInput(input, fast_parse=True) to sniff the data card's prefix, which is itself @args_checked and validates the same input value immediately. jit_parse and problem are likewise validated downstream in every path that actually forwards them (CellModifierInput.__init__ checks problem; every DataClass(...) /DataInput(...) construction checks jit_parse). Confirmed via runtime instrumentation this was another "double guarded" pair found while profiling: ~1002 wasted check_type_and_value calls parsing benchmark/big_model.imcnp. A malformed call now raises from DataInput.__init__ instead of parse_data -- one frame later, same TypeError, no test depended on the exact source function name. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Dispatchers like parse_surface/parse_data need to read one discriminating field (surface type, data prefix) before they know which concrete subclass to build, but had no way to do that without constructing a whole throwaway object first. _JitParser classes (JitSurfParser, JitDataParser) never touch self -- .parse(tokenizer) is a pure function of the token stream -- so this can be done with no instance construction and no @args_checked validation at all. Surface gets a real _BLOCK_TYPE class attribute (previously only ever set as an instance attribute in _init_blank) so the new classmethod can convert a str input to an Input before any instance exists. Not used anywhere yet; wiring up the dispatchers is next. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

parse_surface built a full buffer_surface = Surface(input, jit_parse=True) purely to read _surface_type, discarded it, then built the real concrete subclass separately -- a whole extra @args_checked-validated object per surface just to sniff one field. Confirmed via runtime identity-tracking instrumentation: ('Surface.__init__', 'Surface.__init__') x1003 on benchmark/big_model.imcnp, i.e. Surface.__init__ ran twice for the same input on nearly every surface. Use the new Surface._peek_light_parse to read surface_type with no instance construction on the fast path. On a peek failure, fall back to literally today's original algorithm unchanged (build a real Surface with jit_parse=True, which has its own robust JIT-with-fallback-to-full-parse handling) rather than guessing -- the JIT light parser isn't fully robust and can fail on valid syntax, and silently defaulting to the generic Surface base class on such a failure would misclassify the object, not just run slower. Added test_surface_dispatch_peek_failure_fallback, which monkeypatches the peek to fail and asserts the correct concrete class still comes out. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

parse_data built a full base_input = DataInput(input, fast_parse=True) purely to read .prefix, discarded it, then built the real concrete subclass (e.g. Material) separately. Confirmed via runtime identity-tracking instrumentation: ('DataInput.__init__', 'Material.__init__') x1001 on benchmark/big_model.imcnp. Use the new DataInput._peek_light_parse to read the prefix with no instance construction on the fast path. On a peek failure, fall back to literally today's original algorithm unchanged (build a real DataInput with fast_parse=True, which has its own robust JIT-with-fallback-to-full-parse handling) rather than defaulting the prefix to None -- same rationale as the equivalent Surface fix: the JIT light parser can fail on valid syntax, and guessing "unknown prefix" on such a failure would misclassify the object. Added test_data_parser_peek_failure_fallback (asserts the correct concrete class still comes out) and test_data_parser_malformed_prefix (asserts a genuinely malformed input still raises a sane, contextual error via the fallback construction). Together with the Surface fix, these two changes eliminated both remaining "double guarded" check pairs found while profiling args_checked overhead on this branch; median read_input() time on benchmark/big_model.imcnp dropped from ~1.12s to ~1.02s. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

MicahGale · 2026-07-20T19:04:26Z

Ok, @tjlaboss, Claude was able to find multiple bugs that were causing all objects to be fully parsed on read, hence your weird print behavior. The benchmark is now only do JIT_parse, which be updated to do both. The benchmark now takes < 5 seconds. I also worked with Claude to improve the overhead cost of args_checked.

tjlaboss · 2026-07-20T19:46:07Z

Very good improvement. On a practical model, I observed better than a 2x speedup with the jit_parse branch and the resolution of the issue above.

MicahGale and others added 30 commits November 11, 2025 09:59

Fixed syntax error with finally return in JIT data parser.

317a260

Updated montepy.Universe for new init struct.

1f4faeb

Updated init for importance.

5807c50

Tried UniverseInput init update; redecided infrastructure.

3197f45

Merge branch 'minor-rel-dev' into jit_parse

a6b0f2e

Split parse cell and data tree.

ef3bb78

Fixed blank JIT parsing of importance.

42efcef

Updated init for cell modifiers

26d3f26

Didn't require AST for number.

439b511

Fixed bugs.

625870e

Fixed missing trees decorator for data input.

943c139

Fixed universe bugs.

ed5284a

Handle edge case of void material in pull object.

14cc136

avoid warnings with default cell tree accidentally inserting neutron.

8c1ab38

Removed Update pointers from cell and fixed clone.

a2da554

Updated tests to not use update pointers.

621922e

Worked on fixing cell modifier init. Still very broken.

c05d334

Updated full cell modifier init process to support JIT properly.

e9cf4b5

Add more access to AST.

4e6c27c

Updated tests to provide actually valid syntax tree.

cac3fbe

Updated data tests to force full parsing.

e0bb230

Fixed behavior with cell modifier data shuffling and appending for jit.

17fe924

updated test to be valid.

2347e7d

Prototyped a number impact search method.

ce03a15

Cleaned up parent object finder.

efe4bd6

Mostly implemented push to cells system.

d03003a

Added system to move leading comments around.

084d8f2

Made num object collection tolerant to none numbers.

244b50d

Avoided triggering finding parent objects.

9740289

Made test actually valid

ac8229c

Merge branch 'develop' into jit_parse

b6bb2cd

tjlaboss reviewed Jun 24, 2026

View reviewed changes

Comment thread montepy/mcnp_problem.py

Try to test more parsing branches.

280b944

MicahGale and others added 13 commits July 3, 2026 08:53

Test that problem actually JIT parses.

d336cb1

I can grammar. Fully_parsed not full_parsed.

9ffe00d

Fix Input.search to use re.search instead of re.match.

99c9295

re.match only checks at the start of a string, so patterns would miss keywords anywhere past the first character of a line. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

MicahGale force-pushed the jit_parse branch from 49b5821 to b131981 Compare July 18, 2026 05:45

MicahGale and others added 9 commits July 18, 2026 23:35

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Beta implementation of Just In time parsing#960

Beta implementation of Just In time parsing#960
MicahGale wants to merge 319 commits into
beta_rel_devfrom
jit_parse

MicahGale commented May 29, 2026 •

edited by github-actions Bot

Loading

Uh oh!

tjlaboss commented Jun 24, 2026

Uh oh!

Uh oh!

MicahGale commented Jul 2, 2026

Uh oh!

MicahGale commented Jul 20, 2026

Uh oh!

tjlaboss commented Jul 20, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

MicahGale commented May 29, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pull Request Checklist for MontePy

Description

General Checklist

LLM Disclosure

Documentation Checklist

First-Time Contributor Checklist

Additional Notes for Reviewers

Uh oh!

tjlaboss commented Jun 24, 2026

Uh oh!

Uh oh!

MicahGale commented Jul 2, 2026

Uh oh!

MicahGale commented Jul 20, 2026

Uh oh!

tjlaboss commented Jul 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

MicahGale commented May 29, 2026 •

edited by github-actions Bot

Loading

tjlaboss commented Jul 20, 2026 •

edited

Loading