refactor[next]: Replace apply_to_primitive_constituents with tree_map_type#2626
Open
havogt wants to merge 1 commit into
Open
refactor[next]: Replace apply_to_primitive_constituents with tree_map_type#2626havogt wants to merge 1 commit into
havogt wants to merge 1 commit into
Conversation
…_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.
There was a problem hiding this comment.
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_constituentsand introduced/standardizedtree_map_typeinnext.type_system.type_infoas the canonical API. - Migrated iterator, codegen, and frontend type-deduction call sites to
tree_map_type, including updating custom collection constructors to matchtree_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. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #1570
Description
Removes
type_info.apply_to_primitive_constituentsand migrates all call sites totree_map_type, completing the long-standing TODO inffront/type_info.py.tree_map_type/_tree_map_type_constructorare moved fromffront.type_infodown to the low-leveltype_system.type_info, so iterator- and codegen-layer call sites can use them without an upward dependency onffront;ffront.type_infore-exports them.apply_to_primitive_constituents(which recursed only onts.TupleTypeand treatedts.NamedCollectionTypeas a leaf),tree_map_typerecurses overCOLLECTION_TYPE_SPECSand preserves named-collection structure. This changes behavior only forNamedCollectionTypeinputs, 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_typeis overloaded so its result type reflects the constructor's return (_T | ts.CollectionTypeSpecfor the default,_T | _Cfor a customresult_collection_constructor) instead oftree_map's intentionally weak return type, avoiding casts at the call sites.Pure refactoring, no behavior change at any call site.
Requirements
global_tmps, gtfn codegen, ffront type deduction).