{-# LANGUAGE TypeFamilies, ExplicitForAll, DataKinds, GADTs, MultiParamTypeClasses,
             FlexibleInstances, FlexibleContexts #-}
{-# OPTIONS_GHC -fwarn-incomplete-patterns #-}

module ZipWith where

import Prelude hiding (zipWith)

data Nat = Zero | Succ Nat

-- 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]

-- 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)

-- Use type classes to automatically infer that evidence
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

-- 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 _      _      = []

{-
-- 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

-- 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)
-}

zipWith :: NumArgs numArgs f -> f -> Listify numArgs f
zipWith na f = listApply na (repeat f)

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

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]

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

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