How can Data.Function's `on` function be generalized to work with n-ary functions?

Question

Data.Function in the base package contains a function on :: (b -> b -> c) -> (a -> b) -> a -> a -> c, which is similar to (.) :: (b -> c) -> (a -> b) -> a -> c for unary functions, so I tried to write a function on' :: Int -> ... as a generalization, so that I could write on' 1 length negate, on' 2 length compare, etc., however such a function would not type-check, because the type of the function result of on''s third argument depends on the first argument.

How can I go about writing such a function? Maybe I'd have to wrap functions in a custom data type so that the composed functions' types only depend on the type of the first parameter and the type of the final result?

Answer 1

Here's a possible approach. We start by defining type level naturals.

{-# LANGUAGE ScopedTypeVariables, TypeFamilies, DataKinds, TypeApplications, 
             AllowAmbiguousTypes, MultiParamTypeClasses, FlexibleInstances #-}
{-# OPTIONS -Wall #-}

data Nat = O | S Nat

We define a -> a -> ... a -> b with n arguments.

type family F (n :: Nat) a b where
   F 'O a b = b
   F ('S n) a b = a -> F n a b

Then we introduce a custom class over these naturals for our on, and implement it for every natiral in an inductive way.

class On (n :: Nat) c where
   on :: forall a b. F n b c -> (a -> b) -> F n a c

instance On 'O c where
   on f _g = f

instance On n c => On ('S n) c where
   on f g = \aVal -> on @n @c (f (g aVal)) g

Finally, some examples.

fun2 :: String -> String -> String
fun2 x y = "(" ++ x ++ ", " ++ y ++ ")" 

fun3 :: String -> String -> String -> String
fun3 x y z = "(" ++ x ++ ", " ++ y ++ ", " ++ z ++ ")" 

funG :: Int -> String
funG n = replicate n '#'

test2 :: String
test2 = on @('S ('S 'O)) fun2 funG 1 2

test3 :: String
test3 = on @('S ('S ('S 'O))) fun3 funG 1 2 3

---

A relatively off topic note:

I can't find a way to remove the c argument from the type class. Since c is not determined from the type, it is ambiguous, hence I have to pass it explicitly (either via type application -- as done above -- or a Proxy). However, to pass it, I need c to be in scope. If I remove c from the class it goes out of scope. If I use an instance signature, I can bring c back in scope, but GHC does not recognize it as the same c due to type ambiguity.

OnGeneralization2.hs:18:10: error:
    • Couldn't match type ‘F n a c0’ with ‘F n a c’
      Expected type: F ('S n) b c -> (a -> b) -> F ('S n) a c
        Actual type: F ('S n) b c0 -> (a -> b) -> F ('S n) a c0
      NB: ‘F’ is a type function, and may not be injective
      The type variable ‘c0’ is ambiguous
    • When checking that:
          forall a b c. F ('S n) b c -> (a -> b) -> F ('S n) a c
        is more polymorphic than:
          forall a b c. F ('S n) b c -> (a -> b) -> F ('S n) a c
      When checking that instance signature for ‘on’
        is more general than its signature in the class
        Instance sig: forall a b c.
                      F ('S n) b c -> (a -> b) -> F ('S n) a c
           Class sig: forall a b c.
                      F ('S n) b c -> (a -> b) -> F ('S n) a c
      In the instance declaration for ‘On ('S n)’

Note the last line: they are exactly the same types, but in order to check them for subtyping, GHC still uses fresh Skolem type constants c0 and that makes it fail.

I also tried to make the family injective, but failed.