Variable-arity zipWith

At ICFP in September, an interesting problem was posed: Is it possible to define a variable-arity zipWith in Haskell using GHC 7.6.1? Can we leverage the new expressivity in promoted types and kinds to do away with zipWith3, zipWith4 and friends? The answer turns out to be yes.

Let’s start by enabling a bunch of non-controversial language options and declaring the module:

{-# LANGUAGE TypeFamilies, ExplicitForAll, DataKinds, GADTs,
    	     MultiParamTypeClasses, FlexibleInstances, FlexibleContexts #-}

module ZipWith where

import Prelude hiding (zipWith)

Though promotion is not strictly necessary to pull this off, it turns out to be convenient for GHC to kind-check our code. We define the natural numbers to use at the kind level:

data Nat = Zero | Succ Nat

Now, we need to start thinking about what the type of a variable-arity zipWith must be. Clearly, it will need to take the function to apply and a bunch of lists, but the number of lists is not known when we write the type. We correspondingly don’t know how many arguments the function itself should take. We’ve narrowed our type down to f -> <dragons>, for some function type f. The dragons will have to be some type-level function that evaluates to the correct sequence of arrows and argument types, based on the type substituted for f.

Examples may help here:

• If f is a -> b, then the dragons should be [a] -> [b].
• If f is a -> b -> c, then the dragons should be [a] -> [b] -> [c].
• and so on.

OK. That’s not too hard. We essentially want to map the type-level [] operator over the components of the type of f. However, a problem lurks: what if the final result type is itself an arrow? In the first example above, there is nothing stopping b from being d -> e. This turns out to be a fundemental ambiguity in variable-arity zipWith. Let’s explore this for a moment.

We’ll need a three-argument function to make the discussion interesting. Here is such a function:

splotch :: Int -> Char -> Double -> String
splotch a b c = (show a) ++ (show b) ++ (show c)

Now, there are two conceivable ways to apply splotch with zipWith:

*ZipWith> :t zipWith2 splotch
zipWith2 splotch :: [Int] -> [Char] -> [Double -> String]
*ZipWith> :t zipWith3 splotch
zipWith3 splotch :: [Int] -> [Char] -> [Double] -> [String]

(Here, zipWith2 is really just the zipWith in the Prelude.)

In general, there is no way for an automated system to know which one of these possibilities we want, so it is sensible to have to provide a number to the dragons, which we’ll now name Listify. This number is the number of arguments to the function f. Here is the definition for Listify:

-- Map the type constructor [] over the types of arguments and return value of
-- a function
type family Listify (n :: Nat) (arrows :: *) :: *
type instance Listify (Succ n) (a -> b) = [a] -> Listify n b
type instance Listify Zero a = [a]

Now it would seem we can write the type of zipWith. Except, when we think about it, we realize that the operation of zipWith will have to be different depending on the choice for n in Listify. Because this n is a type, it is not available at runtime. We will need some runtime value that the implementation of zipWith can branch on.

Furthermore, we will need to convince GHC that we’re not doing something very silly, like trying Listify (Succ (Succ (Succ Zero))) (Int -> Int). So, we create a GADT that achieves both of these goals. A value from this GADT will both be a runtime witness controlling how zipWith should behave and will assert at compile time that the argument to Listify is appropriate:

-- Evidence that a function has at least a certain number of arguments
data NumArgs :: Nat -> * -> * where
  NAZero :: NumArgs Zero a
  NASucc :: NumArgs n b -> NumArgs (Succ n) (a -> b)

oneArg = NASucc NAZero
twoArgs = NASucc oneArg
threeArgs = NASucc twoArgs

Finally, we can give the type for zipWith:

zipWith :: NumArgs numArgs f -> f -> Listify numArgs f

Note that, though this zipWith is variable-arity, we still have to tell it the desired arity. More on this point later.

Once we have the type, we still need an implementation, which will need to be recursive both in the length of the lists and the number of arguments. When we think about recursion in the number of arguments to f, currying comes to the rescue… almost. Consider the following:

zipWith threeArgs splotch [1,2] ['a','b'] [3.5,4.5]

We would like a recursive call to come out to be something like

zipWith twoArgs  ['a','b'] [3.5,4.5]

The problem is that there is no replacement for <splotch ??> that works. We want to apply (splotch 1) to the first members of the lists and to apply (splotch 2) to the second members. What we really need is to take a list of functions to apply. Let’s call the function that works with list of functions listApply. Then, the recursive call would look like

listApply twoArgs [splotch 1, splotch 2] ['a','b'] [3.5,4.5]

With such a listApply function, we can now implement zipWith:

zipWith numArgs f = listApply numArgs (repeat f)

The type and implementation of listApply is perhaps a little hard to come up with, but otherwise unsurprising.

-- Variable arity application of a list of functions to lists of arguments
-- with explicit evidence that the number of arguments is valid
listApply :: NumArgs n a -> [a] -> Listify n a
listApply NAZero fs = fs
listApply (NASucc na) fs = listApply na . apply fs
  where apply :: [a -> b] -> [a] -> [b]
        apply (f:fs) (x:xs) = (f x : apply fs xs)
        apply _      _      = []

And now we’re done. Here are some examples of it all working:

example1 = listApply (NASucc NAZero) (repeat not) [False,True]
example2 = listApply (NASucc (NASucc NAZero)) (repeat (+)) [1,3] [4,5]

example3 = zipWith twoArgs (&&) [False, True, False] [True, True, False]
example4 = zipWith twoArgs (+) [1,2,3] [4,5,6]

example5 = zipWith threeArgs splotch [1,2,3] ['a','b','c'] [3.14, 2.1728, 1.01001]

But wait: can we do better? The zipWith built here still needs to be told what its arity should be. Notwithstanding the ambiguity mentioned above, can we somehow infer this arity?

I have not come up with a way to do this in GHC 7.6.1. But, I happen to (independently) be working on an extension to GHC to allow ordering among type family instance equations, just like equations for term-level functions are ordered. GHC will try the first equation and then proceed to other only if the first doesn’t match. The details are beyond the scope of this post (but will hopefully appear later), but you can check out the GHC wiki page on the subject. The following example should hopefully make sense:

-- Count the number of arguments of a function
type family CountArgs (f :: *) :: Nat
type instance where
  CountArgs (a -> b) = Succ (CountArgs b)
  CountArgs result = Zero

This function counts the number of arrows in a function type. Note that this cannot be defined without ordered equations, because there is no way in GHC 7.6.1 to say that result (the type variable in the last equation) is not an arrow.

Now, all we need to do is to be able to make the runtime witness of the argument count implicit through the use of a type class:

-- Use type classes to automatically infer NumArgs
class CNumArgs (numArgs :: Nat) (arrows :: *) where
  getNA :: NumArgs numArgs arrows
instance CNumArgs Zero a where
  getNA = NAZero
instance CNumArgs n b => CNumArgs (Succ n) (a -> b) where
  getNA = NASucc getNA

Here is the new, implicitly specified variable-arity zipWith:

{-# LANGUAGE ScopedTypeVariables #-}
-- Variable arity zipWith, inferring the number of arguments and using
-- implicit evidence of the argument count.
-- Calling this requires having a concrete return type of the function to
-- be applied; if it's abstract, we can't know how many arguments the function
-- has. So, zipWith (+) ... won't work unless (+) is specialized.
zipWith' :: forall f. CNumArgs (CountArgs f) f => f -> Listify (CountArgs f) f
zipWith' f = listApply (getNA :: NumArgs (CountArgs f) f) (repeat f)

This version does compile and work with my enhanced version of GHC. Expect to see ordered type family instances coming to a GHC near you soon!

• GHC 7.6.1 Variable-arity zipWith
• Modified GHC Variable-arity zipWith