Skip to content

Spike: open module-registry seam (illustrates #601) - #615

Draft
jakeb-grant wants to merge 1 commit into
pydantic:mainfrom
jakeb-grant:module-registry-spike
Draft

Spike: open module-registry seam (illustrates #601)#615
jakeb-grant wants to merge 1 commit into
pydantic:mainfrom
jakeb-grant:module-registry-spike

Conversation

@jakeb-grant

@jakeb-grant jakeb-grant commented Jul 23, 2026

Copy link
Copy Markdown

Draft, not for merge. This is a spike to make the idea in #601 concrete,
since you said you were open to seeing what it looks like. The interesting
part is the seam; the module I used to exercise it is throwaway.

One self-contained commit; it builds clean and struct__ops.py passes on both
Monty and CPython.

The idea

Today, adding a native stdlib module means editing the closed StandardLib /
ModuleFunctions enums and threading a new variant through every match that
references them. This replaces that with a name-keyed registry. Building the
seam is a one-time change to the core (the files in the reading order below).
After that, adding a pure-function module like this one is a registration in
two files (modules/mod.rs and registry.rs) plus the module's own file, with
no further core changes. A module that needed a new heap type would touch more,
which is the #568 question below.

The pieces:

  • modules/registry.rs: a static REGISTRY of ModuleDescriptors (plain
    data, no tracker generic) and a registry::call that dispatches by a
    stable ModuleFuncId. It's a plain match, no dyn, so it compiles under
    today's VM<'h, T: ResourceTracker>.
  • Names are interned through the dynamic pool, not StaticStrings, so the
    per-token parse cost stays flat regardless of how many modules exist. And
    prepare interns only the modules a program actually imports, not the whole
    registry, so that cost tracks usage rather than registry size.
  • One immediate Value::RegistryFunction(ModuleFuncId): no heap ref, trivial
    clone/drop, no new HeapData variant. The scalar hot path is untouched,
    which I think is the Add built-in NumPy support #248 regression you flagged.
  • LoadModule now carries a u16 const-pool name id (same shape as
    RaiseImportError). The import gate and load_module resolve registry
    first, then fall back to StandardLib.

What I left open on purpose

About the struct module

Throwaway. calcsize/pack/unpack over a few numeric codes; format and
value errors are ValueError (buffer type errors stay TypeError), no
struct.error/Struct/iter_unpack/native sizes. It's there
only to run the seam end to end and show a module producing and consuming
existing heap types. Please don't review it for parity. Divergences are in
limitations/struct.md.

Reading order

  1. modules/registry.rs — the registry
  2. prepare.rs — lazy per-import name interning
  3. bytecode/compiler.rs compile_import — the "is X in scope" gate
  4. bytecode/vm/mod.rs load_module + vm/call.rs — resolution and dispatch
  5. value.rs — the one new variant

Everything is inside the monty crate (plus the limitations/ docs). No host
crate, wire protocol, or binding changes.

Built with AI assistance (noted in the commit trailer). What I'm after is your
read on whether this is the shape you had in mind, not a merge review.


Summary by cubic

Introduces an open, name-keyed stdlib module registry so new native modules can be added via registration instead of editing closed enums; adds an illustrative struct module to exercise the seam end to end. Implements the extensibility direction discussed in #601 while keeping hot paths unchanged.

  • New Features

    • Added modules/registry.rs with ModuleDescriptor, ModuleFuncId, REGISTRY, lookup/is_registered, create_module, call (append-only dispatch), plus descriptor_names and function_name.
    • New Value::RegistryFunction with VM call dispatch and repr; identity/serde and hashing support added.
    • Compiler emits LoadModule with a u16 const index for the module name; the import gate checks the registry first, then StandardLib.
    • prepare lazily interns names only for imported registry modules (walks nested imports).
    • Added minimal struct module (calcsize, pack, unpack) plus tests and docs under limitations.
  • Refactors

    • load_module now resolves by interned name (registry-first, fallback to StandardLib); VM reads the name from the const pool at runtime.
    • Opcode LoadModule now uses a u16 const-pool operand; stack effects updated accordingly.
    • VM call path dispatches RegistryFunction via registry::call; module attribute calls use vm.call_function.
    • Added ExcType::runtime_error for invariant failures like invalid registry ids from snapshots; REPL/callability checks include RegistryFunction.

Written for commit d1b2fa9. Summary will update on new commits.

Review in cubic

@codecov

codecov Bot commented Jul 23, 2026

Copy link
Copy Markdown

@codspeed-hq

codspeed-hq Bot commented Jul 23, 2026

Copy link
Copy Markdown

Merging this PR will not alter performance

✅ 26 untouched benchmarks
⏩ 16 skipped benchmarks1


Comparing jakeb-grant:module-registry-spike (d1b2fa9) with main (0a5fdb7)

Open in CodSpeed

Footnotes

  1. 16 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports.

@samuelcolvin

Copy link
Copy Markdown
Member

I'm afk right now, so haven't looked at the code yet, but wouldn't it be more elegant long term to add a Dynamic(dyn PyTrait) variant to the heap data enum, then use that?

@jakeb-grant

Copy link
Copy Markdown
Author

Agreed that's a nicer end state, and I think in line with #568. I left it out here because dyn PyTrait needs the resource-tracker generic gone first. So the spike doesn't depend on it and is built to become a dyn PyTrait table once #613 lands. Happy to table till then or redo on top of #613 if that works.

Spike for the extensibility direction discussed on pydantic#601: replace the
closed `StandardLib`/`ModuleFunctions` enum path with a name-keyed,
append-only registry so adding a native stdlib module is a registration,
not a core edit.

The seam:
- `modules/registry.rs`: a `static REGISTRY` of `ModuleDescriptor`s and a
  free fn `registry::call` (dispatch by stable `ModuleFuncId`, a plain
  match — the simplest shape; now that pydantic#613 has removed the VM's
  `ResourceTracker` generic, a fn-pointer table or `dyn` registration
  could replace the match without touching callers).
- Dynamic interning: module/function names go through the dynamic string
  pool, not `StaticStrings`, so the per-token parse cost stays flat in the
  number of registered modules; `prepare` interns only the modules a program
  actually imports (full-tree walk, so nested `import`s are covered), so that
  cost tracks usage rather than registry size.
- One immediate `Value::RegistryFunction(ModuleFuncId)` variant, no heap
  reference, trivial clone/drop — the scalar hot path is untouched and no
  `HeapData` variant is added (the pydantic#248 regression concern).
- `LoadModule` now carries a u16 const-pool module-name id (mirroring
  `RaiseImportError`); the compiler import gate and `load_module` resolve
  registry-first, then fall back to `StandardLib`.

`struct` (`calcsize`/`pack`/`unpack` over standard-size numeric codes) is an
illustrative stub to exercise the seam end to end without a new heap type;
it is not a full `struct` implementation. Resource use is charged before
every format-derived allocation (`pack` buffer, `unpack` output tuple, and
the parsed-field vector), and huge repeat counts are rejected rather than
overflowing. Everything is contained to the `monty` crate. See
`limitations/struct.md`.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@jakeb-grant
jakeb-grant force-pushed the module-registry-spike branch from 67d6e89 to d1b2fa9 Compare July 24, 2026 16:27
@jakeb-grant

Copy link
Copy Markdown
Author

Rebased onto main now that #613 is in. The impl ResourceTracker annotations dropped out of the registry.

@samuelcolvin on Dynamic(dyn PyTrait): with #613 merged, the remaining hang up is object safety. #568 moves the binary operators to &Value (py_add_impl etc.), but py_cmp still takes &Self and py_repr_fmt takes &mut impl Write, so those two would need the same treatment before the trait is dyn-compatible. Once done, the registry here doesn't change shape, only what a ModuleFuncId resolves to does (a table of trait objects instead of the match).

Happy to take a pass at the dyn PyTrait version of this once #568 hits.

@samuelcolvin

Copy link
Copy Markdown
Member

Makes sense. In general I think the direction is good. I'll let @davidhewitt review the implementation.

A few notes:

  • if we're actually adding the struct module, we should add it to typeshed
  • should we make all changes here, or merge this with two completing approaches, then remove the other approach later
  • I think this is moving in the right direction, but just for your awareness: I'd like to add support for loading modules from the filesytem (e.g. the monty-fs machinery) in future
My claude conversation in case it's helpful

PR #615 — open module-registry seam: review conversation

Notes from reviewing d1b2fa9f ("Add open module-registry seam (illustrative
struct module)") against main, in the context of
issue #601.


Q: Compare this PR to main, and look at issue #601. Explain simply how it works — I'm not clear how it makes it easier to add new modules?

What the PR is

Issue #601 is a direction check, not a feature request: jakeb-grant asked whether
we'd accept a first-party seam for native modules, in two parts — (a) an open
module registry, (b) a HeapData::Foreign variant later. David said "open to
seeing what the registry looks like". This branch is the spike for part (a) only,
with struct as a throwaway proof module.

How it works

On main, an import is resolved entirely at compile time against a closed enum:

  • StandardLib::from_string_id(name) — which only works because every module name
    (math, re, …) is a variant of StaticStrings, the 362-variant enum in
    intern.rs.
  • The compiler emits LoadModule with a u8 operand = the StandardLib
    discriminant
    .
  • At runtime StandardLib::from_repr(id).create(vm) builds the module, setting
    attributes keyed by more StaticStrings variants (Sqrt, Dumps, …), each
    holding a Value::ModuleFunction(ModuleFunctions::Math(MathFunctions::Sqrt))
    a nested enum, one layer per module.

On this branch, that lookup becomes name-based and late:

  • LoadModule's operand becomes a u16 const-pool index pointing at the module
    name string
    (the same encoding RaiseImportError already used).
  • At runtime load_module takes the name, tries registry::lookup first, falls
    back to StandardLib.
  • The registry is a static REGISTRY: &[ModuleDescriptor] — each descriptor is
    just { name: "struct", functions: &[("calcsize", ModuleFuncId(0)), …] }, plain
    &'static strs, no StaticStrings.
  • New Value::RegistryFunction(ModuleFuncId), dispatched by registry::call's
    match id.0.
  • prepare.rs walks the AST for import statements and dynamically interns the
    names of just the registered modules that are actually imported, so runtime can
    look them up.

So what actually gets easier?

Concretely, adding a module today touches four places:

  1. intern.rs — one StaticStrings variant for the module name plus one per
    function name
  2. modules/mod.rsStandardLib variant + from_string_id arm + create arm
  3. modules/mod.rs again — ModuleFunctions variant + Display arm + call arm
  4. the module's own file — its private XxxFunctions enum, Display,
    create_module, call

With the registry it's two:

  1. the module's own file — a DESCRIPTOR const + the functions
  2. registry.rs — one entry in REGISTRY, one dispatch arm per function

The real payoff isn't the edit count, it's step 1. StaticStrings derives
strum's EnumString, so from_str is a generated match that every
multi-character token probes once at parse time. Adding numpy's ~200 names means
+200 arms in a match every program pays for, whether or not it imports numpy.
That's the startup regression that sank #248. The registry keeps module vocabulary
out of that match entirely — costs are paid only by programs that import the
module. That argument holds.

Where I'd push back

The edit-count win is modest and partly self-inflicted. ModuleFuncId is a
flat, hand-assigned id space shared across all modules, and the descriptor's
ids must be manually kept in sync with the arm numbers in registry::call
nothing checks it. That's strictly worse than the nested
ModuleFunctions::Math(MathFunctions::Sqrt) enums it replaces, where each module
owns its own namespace and the compiler enforces the pairing. function_name()
also does a linear scan over every function in every module to render a repr.

The stated reason for ids-not-fn-pointers ("plain data, so it lives in a
static") is wrong as written — fn pointers are fine in a static. The actual
reason is that Value::RegistryFunction is serde-serialized into snapshots and
folded into id(), so it needs a stable number. That's legitimate, but the number
could be derived from (registry index, function index) automatically instead of
hand-written twice.

Also worth checking: create_module expect()s that names were pre-interned by
prepare, and prepare only seeds modules it finds via a hand-written AST walk over
Import/ImportFrom nodes. Any path that reaches a module without that walk
having seen it — a REPL feed importing a module a prior feed didn't, or a node kind
the walk's _ => {} arm misses — is a panic rather than an error.


Q: We want to support modules loaded from the monty filesystem, e.g. from the monty-fs machinery for reading files. Will this support that, or should that come later?

Later — but this PR moves one step in the right direction.

The half that generalizes is the opcode change. Going from LoadModule <u8 StandardLib discriminant> to LoadModule <const-pool module name> is exactly what
a filesystem import needs: the VM ends up holding a name and resolving it at
runtime, so load_module becomes the natural place for a third fallback after
registry and StandardLib. Keep that regardless.

The half that doesn't generalize is the interning strategy, and that's where the
real blocker sits:

Interns is immutable at runtime. get_string_id_by_name is &self; only
InternerBuilder has intern. Growing the table means Interns → InternerBuilder::from_interns → prepare/compile → Interns, which is what
repl.rs does between feeds — outside the VM run loop. You cannot intern a new
string from inside LoadModule. A filesystem module's source contains arbitrary
identifiers that by definition weren't seen at prepare time, so this PR's approach
(walk the AST at prepare time, seed the names, then expect("must be pre-interned") at runtime) works for a fixed static registry and is a dead end for
anything loaded from disk.

Beyond interning, a monty-fs import needs three more things this PR doesn't touch:

  • Suspension for the read. The interpreter does no I/O, so import foo must
    suspend with an OsFunctionCall, unwind to the host, and resume mid-opcode with
    the source — a suspend point in the middle of LoadModule, which currently
    doesn't exist.
  • Compile at runtime, then execute a nested frame to build the namespace,
    rather than the current "allocate a Module and set attributes" shape.
  • A module cache. Today create_module allocates a fresh Module on every
    LoadModule (true on main too, for StandardLib). Tolerable for pure-function
    stdlib modules; not tolerable once importing has side effects and identity, where
    you need sys.modules semantics.

Also, the compiler's static gate (is_registered(name) || StandardLib::from_string_id(name), else emit RaiseImportError) has to go away
for filesystem imports — availability isn't knowable at compile time.

So: it's a genuinely bigger piece of work than the registry, and none of it is
blocked by this PR. Suggestion if we take this branch: keep the name-carrying
LoadModule, but don't bless the prepare-time seeding as the mechanism — flag it
as spike-only, because part 2 will have to replace it.


Q: Is the search for a module linear? Should we use a hashmap or similar to make lookup faster?

Not worth a HashMap, but function_name should change.

Where the scans are and how hot:

  • lookup_by_name — linear over REGISTRY, called once per import statement at
    compile time and once per LoadModule execution. Imports are rare and modules
    will number in the tens. At that size a scan of short &strs beats a HashMap,
    which has to hash the whole key before it can do anything; you'd be adding a
    structure to own and build for no measurable win. If it ever does matter the
    answer is a compile-time perfect hash (phf) over a static, not a runtime map.
  • function_name(id) — this one is REGISTRY.iter().flat_map(|m| m.functions).find(...), i.e. O(all functions across all modules), and it's on
    the repr path. Still cold, but it scales with the wrong quantity.

The fix worth making isn't a map at all: make ModuleFuncId a (module_index, function_index) pair instead of a flat hand-assigned u16. That turns
function_name into two array indexes, and it simultaneously kills the hazard
noted above — that the ids in each DESCRIPTOR and the arm numbers in
registry::call's match are hand-written in two places with nothing checking
they agree.

The hashmap that will be genuinely needed is a different one: the per-session
sys.modules cache, keyed by name at runtime. That's part of the filesystem work,
not the static registry.


Q: What is phf?

phf is a Rust crate — "perfect hash function". It builds a lookup table at
compile time
whose hash function is chosen specifically so that the keys you gave
it collide zero times:

static MODULES: phf::Map<&'static str, &ModuleDescriptor> = phf_map! {
    "struct" => &struct_::DESCRIPTOR,
    "array"  => &array::DESCRIPTOR,
};

The difference from HashMap:

  • HashMap is built at runtime, lives on the heap, has to be constructed and
    owned somewhere, and a lookup means hash → bucket → possibly walk a collision
    chain → compare keys.
  • phf::Map is a static in the binary — no allocation, no construction, no
    ownership question. A lookup is hash → one array index → one key comparison to
    confirm. No chain, because there are no collisions by construction.

It's the standard tool for "fixed set of string keys known at compile time", which
is exactly what a module registry is. Keyword tables in parsers and lexers use it a
lot.

That said, for this case I'd not bother. phf isn't currently a dependency of the
workspace, and you'd be adding it plus a proc-macro build step to speed up a lookup
that happens once per import statement across maybe 20 keys — where a linear scan
of short &strs is already competitive, since hashing the query string costs about
as much as a handful of short comparisons.

It came up mainly to make the point that if the scan ever becomes a real cost,
the fix is a compile-time table rather than a runtime HashMap — the latter would
be strictly worse here. The scan that actually deserves attention is
function_name's O(all functions), and that's fixed by changing the id
representation, not by adding any map.

@jakeb-grant

Copy link
Copy Markdown
Author
  • typeshed: agreed. Happy to add that if this moves past a spike.
  • The compiler's import gate and load_module both try the registry first and fall back to StandardLib, so built-ins can migrate one descriptor at a time and drop the enums later. So I'd probably remove the other approach later.
  • Per monty-fs, I agree with the review. The name-carrying LoadModule is the piece that generalizes; the prepare-time seeding is not, since interns are immutable at runtime (intern only exists on InternerBuilder, rebuilt between feeds). A disk import needs runtime interning, a suspend point inside LoadModule, and a module cache with sys.modules semantics. I don't think any of that is blocked by this PR, and I'd read the seeding as registry-only.
A few things I had Fable 5 vet for me
  • Flat hand-assigned ModuleFuncId: agreed. Deriving it from (module index, function index) removes the hand-written id copy and makes function_name two array indexes. Dispatch would stay a hand-ordered match unless the table also carries fn pointers, which per the next point it can. Either way the id remains an append-only contract, since it's serialized into dumps and folded into id().
  • "No fn pointers, so it lives in a static": right, that's wrong as written — a fn-pointer table is fine in a static. The actual reason for ids is that they're serialized into snapshots and identity, which fn pointers can't be. Should fix the docstring.
  • The REPL panic doesn't exist: every feed (REPL or one-shot) goes through prepare_with_existing_names, which seeds from that feed's own AST on top of the carried-forward intern table, and interns commit even when a feed raises. I tested the exact scenario (second feed imports struct when the first didn't), plus imports nested in for/while/if/try and an import executed by a later feed's call: no panic. The walk covers every body-carrying node variant, match is rejected at parse, and there's no __import__, so an import can't appear anywhere the walk doesn't visit.
  • The kernel that is real: a hand-crafted/corrupt dump (interns containing struct but not calcsize) does reach the expect in create_module. No legitimate path does, and corrupt dumps can already panic in several pre-existing places, but it's inconsistent with registry::call in the same file, which degrades the same case to a catchable error. Will swap the expect to match.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants