Skip to content

refactor[next]: Replace apply_to_primitive_constituents with tree_map_type#2626

Open
havogt wants to merge 1 commit into
GridTools:mainfrom
havogt:migrate-apply-to-primitive-constituents
Open

refactor[next]: Replace apply_to_primitive_constituents with tree_map_type#2626
havogt wants to merge 1 commit into
GridTools:mainfrom
havogt:migrate-apply-to-primitive-constituents

Conversation

@havogt
Copy link
Copy Markdown
Contributor

@havogt havogt commented Jun 5, 2026

Closes #1570

Description

Removes type_info.apply_to_primitive_constituents and migrates all call sites to tree_map_type, completing the long-standing TODO in ffront/type_info.py.

  • The canonical tree_map_type / _tree_map_type_constructor are moved from ffront.type_info down to the low-level type_system.type_info, so iterator- and codegen-layer call sites can use them without an upward dependency on ffront; ffront.type_info re-exports them.
  • Unlike apply_to_primitive_constituents (which recursed only on ts.TupleType and treated ts.NamedCollectionType as a leaf), tree_map_type recurses over COLLECTION_TYPE_SPECS and preserves named-collection structure. This changes behavior only for NamedCollectionType inputs, which are lowered to plain tuples (named_collections_to_tuple_types) before any iterator-/codegen-stage call site — so the migrated sites are unaffected.
  • tree_map_type is overloaded so its result type reflects the constructor's return (_T | ts.CollectionTypeSpec for the default, _T | _C for a custom result_collection_constructor) instead of tree_map's intentionally weak return type, avoiding casts at the call sites.

Pure refactoring, no behavior change at any call site.

Requirements

  • All fixes and/or new features come with corresponding tests. — behavior-preserving refactor, covered by existing suites (type system, iterator type inference, global_tmps, gtfn codegen, ffront type deduction).
  • Important design decisions have been documented in the appropriate ADR — N/A, no user-facing design change.

…_type

Remove `type_info.apply_to_primitive_constituents` and migrate all call
sites to `tree_map_type` (a `utils.tree_map`-based helper over
`COLLECTION_TYPE_SPECS`). The canonical `tree_map_type` /
`_tree_map_type_constructor` now live in the low-level
`type_system.type_info` so every caller layer can reach them;
`ffront.type_info` re-exports them.

Unlike the old function, which recursed only on `ts.TupleType` and
treated `ts.NamedCollectionType` as a leaf, `tree_map_type` recurses into
named collections and preserves their structure.

`tree_map_type` is overloaded so its result type reflects the
constructor's return (`_T | ts.CollectionTypeSpec` for the default,
`_T | _C` for a custom one) instead of `tree_map`'s intentionally weak
return type, avoiding casts at most call sites.
Copy link
Copy Markdown

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR completes the refactor away from type_info.apply_to_primitive_constituents by removing it and migrating all call sites to the more general tree_map_type, while relocating the canonical tree_map_type implementation to the lower-level next.type_system.type_info module (and re-exporting it from ffront.type_info) to avoid upward dependencies.

Changes:

  • Removed apply_to_primitive_constituents and introduced/standardized tree_map_type in next.type_system.type_info as the canonical API.
  • Migrated iterator, codegen, and frontend type-deduction call sites to tree_map_type, including updating custom collection constructors to match tree_map’s (value, elts) constructor convention.
  • Tightened/clarified typing at a few call sites (e.g., explicit assertions for tuple/field results, adjusted return type annotation in type_synthesizer).

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated no comments.

Show a summary per file
File Description
src/gt4py/next/type_system/type_info.py Removes apply_to_primitive_constituents; adds canonical _tree_map_type_constructor + tree_map_type wrapper around next.utils.tree_map.
src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py Switches tuple-leaf processing to tree_map_type and updates tuple constructors to the (value, elts) convention.
src/gt4py/next/iterator/type_system/type_synthesizer.py Replaces dtype/field construction mappings with tree_map_type; updates return typing and adds sanity assertions.
src/gt4py/next/iterator/type_system/inference.py Updates temporary type construction to use tree_map_type and asserts expected mapped result shape.
src/gt4py/next/iterator/transforms/global_tmps.py Migrates tuple-structured name/dtype/domain propagation logic to tree_map_type with an explicit tuple constructor.
src/gt4py/next/iterator/transforms/fuse_as_fieldop.py Replaces dtype extraction mapping with tree_map_type for list-type checks.
src/gt4py/next/ffront/type_info.py Re-exports canonical _tree_map_type_constructor and tree_map_type from next.type_system.type_info.
src/gt4py/next/ffront/foast_passes/type_deduction.py Migrates astype return type mapping to tree_map_type and removes an unnecessary type-ignore on where return-type deduction.

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