Investigate Lambda Set Specialisation

[athas]

Futhark imposes certain draconian type system restrictions in order to guarantee efficient defunctionalisation. Essentially we guarantee that any functional value can have only a single possible form, and then we specialise every higher-order function.

The paper Better Defunctionalization through Lambda Set Specialization shows a defunctionalisation technique that allows functional values to come from multiple lambdas, while precisely tracking exactly which are possible.

We should consider reimplementing defunctionalisation in terms of Lambda Set Specialization. Returning functions from branches can certainly result in slower code as the representation becomes a sum type, but it would be opt-in. Any code written under our current restrictions would be unaffected.

Specifically, after implementing Lambda Set Specialization we could allow functions to be returned from branches and used as loop variables. We still would not allow functions in arrays, as this can induce irregular arrays. However, it would mean that size-lifted and lifted types would have the same restrictions (can't be put in arrays) and hence could be combined into a single notion of lifted types.

[athas]

Apart from the paper, maybe some inspiration can be derived from the implementation at https://github.com/morphic-lang/morphic/tree/main/src

I think it's in these files:

https://github.com/morphic-lang/morphic/blob/main/src/annot_closures.rs

https://github.com/morphic-lang/morphic/blob/main/src/closure_specialize.rs