google-research / dex-lang

Research language for array processing in the Haskell/ML family
BSD 3-Clause "New" or "Revised" License
1.56k stars 106 forks source link

Implement traversals over the type/atom/lambda/pi components of CoreIR/SimpIR terms #1288

Closed dougalm closed 1 year ago

dougalm commented 1 year ago

The current GenericTraversal machinery goes under binders generically so it has to make assumptions about the name environment and how to extend it. This makes it less applicable than I'd initially hoped. For example, the DCE and Occ passes can't use it because they're bottom-up instead of top-down.

This patch adds a name-agnostic system that lets you traverse the components of a term in an aribtrary (e.g. non-name-carrying) monad. The term's components are a mixture of non-atom names, atoms, types, lambda terms and pi terms. The caller-supplied functions for handling lambda and pi presumably go under their binders, but this only needs to work in the caller's specific monad.

This machinery is somewhere between GenericE and GenericTraversal in terms of how much is assumes about the IR structure and the monad you're working in.