google-research / dex-lang

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

Reorganize primitive ops #1286

Closed dougalm closed 1 year ago

dougalm commented 1 year ago

This is in preparation for treating types as distinct from atoms in SimpIR and then allowing expressions in types. I switched over the remaining generic "functor" primitives to ADTs that explicitly contain atoms.

The challenge is that we can no longer derive instances for generic traversals like Traversable and we can't use the unit-leaf (e.g. MiscOp ()) versions to represent op "names". So this change also introduces some new machinery for generically handling primitive ops.

Our main representation of primitive operations constrains their arity and the sorts of arguments they take, with constructors like RunReader (Atom r n) (LamExpr r n) which is convenient when we're pattern-matching on a particular primitive. But in other places we want a more generic representation - when we're parsing, printing and generically traversing for substitution or whatever.

This patch introduces a class GenericOp which lets you describe an isomorphism between and op and an (opConstant, [LamExpr], [Atom]) triple (soon to be extended with type args too). The opConstant field is an associated type and it should have instances for printing and parsing.