Nary map in Idris

nary.idr

import Data.List
import Data.Vect
  
total
repeatAndRet : (n : Nat) -> Type -> Type -> Type
repeatAndRet Z argType resultType = resultType
repeatAndRet (S n) argType resultType = argType -> repeatAndRet n argType resultType
  
total
nary : Nat -> Type-> Type
nary n ty = repeatAndRet n ty ty

infixl 9 >.<
total
(>.<) : (a -> b) -> (repeatAndRet n t a) -> repeatAndRet n t b
(>.<) {n=      Z  } q v = q v
(>.<) {n=   (S Z) } q f = q . f
(>.<) {n=(S (S n))} q f = \y => ((>.<) {n=(S n)}) q (f y)
  
total
naryCons : (n : Nat) -> repeatAndRet n t (Vect n t)
naryCons Z = []
naryCons (S Z) = (::[])
naryCons (S n) = \x => with Vect ((::) x) >.< (naryCons n)
  
total
applyN : {n : Nat} -> repeatAndRet n t o -> Vect n t -> o
applyN {n=Z} fun [] = fun
applyN {n=(S n)} fun (h::t) = applyN (fun h) t

total
collect : Monad m => (a -> m b) -> Vect n a -> m $ Vect n b
collect {b=b} _ [] = pure $ the (Vect 0 b) []
collect fun (h::t) = liftA (::) (fun h) <*> collect fun t

total
collectTails : Vect n (List t) -> Maybe $ Vect n (List t)
collectTails = collect tail'

total
collectHeads : Vect n (List t) -> Maybe $ Vect n t
collectHeads = collect head'

zipVects : { n : Nat } -> repeatAndRet n t o -> Vect n (List t) -> List o
zipVects v [] = [v]
zipVects fun vs = case (collectHeads vs) of
  Nothing => []
  Just hs => (applyN fun hs) :: (maybe [] (zipVects fun) $ collectTails vs)

map' : (n : Nat) -> (nary n t) -> nary n (List t)
map' n fun = (zipVects fun) >.< (naryCons n)