Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Variable-arity zipWith in terms of Applicative ZipList

I was implementing "variable-arity zipWith" from Richard Eisenberg's thesis (recommended) when I noticed it used

apply :: [a -> b] -> [a] -> [b]
apply (f:fs) (a:as) = f a : apply fs as
apply _      _      = []

and

repeat :: a -> [a]

which correspond to

(<*>) @ZipList :: ZipList (a -> b) -> ZipList a -> ZipList b
pure  @ZipList :: a -> ZipList a

So I decided to implement it parameterised over Applicative (don't take the formatting seriously, I'm going through a phase)

{-# Language RankNTypes            #-}
{-# Language LambdaCase            #-}
{-# Language TypeApplications      #-}
{-# Language ScopedTypeVariables   #-}
{-# Language DataKinds             #-}
{-# Language KindSignatures        #-}
{-# Language GADTs                 #-}
{-# Language TypeFamilies          #-}
{-# Language FlexibleInstances     #-}
{-# Language MultiParamTypeClasses #-}
{-# Language UndecidableInstances  #-}
{-# Language FlexibleContexts      #-}
{-# Language AllowAmbiguousTypes   #-}
{-# Language InstanceSigs          #-}

import Unsafe.Coerce
import Control.Applicative
import Data.Kind

data N = O | S N

data
 NumArgs :: N -> Type -> Type
 where
 NumArgsO :: NumArgs   (O)      (a)
 NumArgsS :: forall a b n.
             NumArgs   (n)      (b)
          -> NumArgs (S n) (a -> b)

-- CNumArgs :: N -> Type -> Constraint
class CNumArgs n a where
 getNumArgs :: NumArgs n a
instance CNumArgs O a where
 getNumArgs :: NumArgs O a
 getNumArgs = NumArgsO
instance (CNumArgs n b, a_b ~ (a -> b)) => CNumArgs (S n) a_b where
 getNumArgs :: NumArgs (S n) (a -> b)
 getNumArgs = NumArgsS getNumArgs

-- Listify :: N -> (Type -> Type) -> (Type -> Type)
type family
 Listify (n :: N) (f :: Type -> Type) (a :: Type) :: Type where
 Listify O     f a        = f a
 Listify (S n) f (a -> b) = f a -> Listify(n) f b

listApply :: Applicative f
          => f a
          -> NumArgs(n) a
          -> Listify(n) f a
listApply fs = \case
 NumArgsO ->
  fs
 NumArgsS numArg -> \arg ->
  listApply (fs <*> arg) numArg

-- NotFun :: Type -> Constraint
class    CountArgs a ~ O => NotFun a
instance CountArgs a ~ O => NotFun a

-- CountArgs :: Type -> N
type family
 CountArgs (a :: Type) :: N where
 CountArgs (_ -> b) = S (CountArgs b)
 CountArgs _        = O

-- Arity :: N -> Type -> Constraint
class    (CNumArgs(n) a, n ~ CountArgs a) => Arity(n) a
instance (CNumArgs(n) a, n ~ CountArgs a) => Arity(n) a

gzip :: forall f n a.
        Applicative f
     => Arity(n) a
     => a
     -> Listify(n) f a
gzip fs = listApply @f (pure fs) (getNumArgs @n @a)

gzip_ :: forall n a.
         Arity(n) a
      => a
      -> Listify(n) [] a
gzip_ = unsafeCoerce (gzip @ZipList @(n) @a)

one :: NotFun a => a -> [a]
one = gzip_

two :: NotFun b => (a -> b) -> ([a] -> [b])
two = gzip_

three :: NotFun c => (a -> b -> c) -> ([a] -> [b] -> [c])
three = gzip_

We want the behaviour of Applicative ZipList but in the end we want [] instead of ZipList, we know that [a] and ZipList a are Coercible meaning that they share a memory representation (can be reified with Coercion)

coerce :: [a] -> ZipList a

Coercion :: Coercion [] ZipList

but GHC isn't smart enough to reason that the type function Listify ZipList (CountArgs a) a has the same representation in memory as Listify [] (CountArgs a) a so I use unsafeCoerce, feel free to post ways to get around it

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment