How to make Haskell compute the correct polymorphic type?

Question

I just realized how useful the little on-function can be.

Ex:

orderByLength = sortBy (compare `on` length) 

But unfortunately, the inferred types can be somewhat counter-intuitive.

According to the very definition

f `on` g = \x y -> f (g x) (g y)

one could e.g. replace

(==) `on` length

with

\x y -> (length x) == (length y)

But both have different types!

The first has [a] -> [a] -> Bool whereas the second has the correct, more generic type of [a] -> [b] -> Bool.

This disallows obviously correct terms like (on (==) length) [1, 2, 3] ["a", "b", "c"] (which should yield True but now even fails type-checking).

I know this restriction comes up due to the usage of first-rank types, but how to overcome this? Can someone formulate an implementation of on that can deal correctly with polymorphic functions (using universal quantification/rank-n types)?

Answer 1

{-# LANGUAGE Rank2Types #-}
on' :: (a -> a -> b) -> (forall d. c d -> a) -> c e -> c f -> b
on' f g x y = f (g x) (g y)

This results in

Prelude> :t on' (==)
on' (==) :: (Eq a) => (forall d. c d -> a) -> c e -> c f -> Bool
Prelude> :t on' (==) length
on' (==) length :: [e] -> [f] -> Bool

On the other hand, this signature also makes flip on' id illegal, which is somewhat less than desirable.

---

{-# LANGUAGE TemplateHaskell #-}
import Language.Haskell.TH
onE f g = do
    x <- newName "x"
    y <- newName "y"
    lamE [varP x, varP y] $ f `appE` (g `appE` varE x) `appE` (g `appE` varE y)

Prelude> :set -XTemplateHaskell
Prelude> $(onE [|(==)|] [|length|]) [1,2,3] ["a","b","c"]
True
Prelude> $(onE [|(==)|] [|id|]) 4 5
False