[Haskell-cafe] Composition of n-arity functions

[Bas van Dijk]

Hello,

In the levmar binding[1][2] me and my brother are working on, I need a
function composition operator that is overloaded to work on functions
of any arity. Basically its type needs to be something like the
following:

(.*) :: (b -> c) -> NFunction n a b -> NFunction n a c

where 'NFunction n a b' represents the function 'a_0 -> a_1 -> ... -> a_n -> b'

I have written the following implementation:

(my question to this list is below)

-- Begin -----------------------------------------------------------------------

{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE FlexibleContexts #-}

module ComposeN where

--------------------------------------------------------------------------------
-- Type-level naturals (from an idea by Ryan Ingram)
--------------------------------------------------------------------------------

data Z = Z
newtype S n = S n

type N0 = Z
type N1 = S N0
type N2 = S N1
type N3 = S N2

class Nat n where
   caseNat :: forall r.
              n
           -> (n ~ Z => r)
           -> (forall p. (n ~ S p, Nat p) => p -> r)
           -> r

instance Nat Z where
   caseNat _ z _ = z

instance Nat n => Nat (S n) where
   caseNat (S n) _ s = s n

induction :: forall p n. Nat n
          => n
          -> p Z
          -> (forall x. Nat x => p x -> p (S x))
          -> p n
induction n z s = caseNat n isZ isS
    where
      isZ :: n ~ Z => p n
      isZ = z

      isS :: forall x. (n ~ S x, Nat x) => x -> p n
      isS x = s (induction x z s)

newtype Witness x = Witness { unWitness :: x }

witnessNat :: forall n. Nat n => n
witnessNat = theWitness
    where
      theWitness = unWitness $ induction (undefined `asTypeOf` theWitness)
                                         (Witness Z)
                                         (Witness . S . unWitness)


--------------------------------------------------------------------------------
-- N-arity functions
--------------------------------------------------------------------------------

-- | A @NFunction n a b@ is a function which takes @n@ arguments of
-- type @a@ and returns a @b at .
-- For example: NFunction (S (S (S Z))) a b ~ (a -> a -> a -> b)
type family NFunction n a b :: *

type instance NFunction Z     a b = b
type instance NFunction (S n) a b = a -> NFunction n a b

-- | @f .* g@ composes @f@ with the /n/-arity function @g at .
(.*) :: forall n a b c. (ComposeN n) => (b -> c) -> NFunction n a b ->
NFunction n a c
(.*) = compose (witnessNat :: n) (undefined :: a)

infixr 9 .* -- same as .

class Nat n => ComposeN n where
    compose :: forall a b c. n -> a ->
               (b -> c) -> NFunction n a b -> NFunction n a c

-- Note that the 'n' and 'a' arguments to 'compose' are needed so that the type
-- checker has enough information to select the right 'compose' instance.

instance ComposeN Z where
    compose Z _ = ($)

instance ComposeN n => ComposeN (S n) where
    compose (S n) (_ :: a) f g = compose n (undefined :: a) f . g


--------------------------------------------------------------------------------
-- Test
--------------------------------------------------------------------------------

foo :: NFunction N3 Integer Integer
foo x y z = x + y + z

bar :: Integer -> Integer
bar k = k - 1

test1 = compose (witnessNat :: N3)
                (undefined  :: Integer)
                bar foo 1 2 3

test2 = (bar .* foo) 1 2 3


-- The End ---------------------------------------------------------------------

The problem is test1 type checks and evaluates to 5 as expected but
test2 gives the following type error:

Couldn't match expected type `NFunction n a Integer'
           against inferred type `Integer -> Integer -> Integer -> Integer'
    In the second argument of `(.*)', namely `foo'

However if I ask ghci to infer the type of (bar .* foo) I get:

*ComposeN>:t (bar .* foo)
(bar .* foo)
  :: (Integer -> Integer -> Integer -> Integer
        ~
      NFunction n a Integer,
      ComposeN n) =>
     NFunction n a Integer

Here we see that the context contains the type equality:

(Integer -> Integer -> Integer -> Integer ~ NFunction n a Integer

So why is ghci unable to match the expected type `NFunction n a Integer'
against the inferred type `Integer -> Integer -> Integer -> Integer'
while the context contains just this equality?

regards,

Bas

[1] http://code.haskell.org/~basvandijk/code/bindings-levmar/
[2] http://code.haskell.org/~basvandijk/code/levmar/