Idris implementation of nmap from https://www.seas.upenn.edu/~sweirich/papers/aritygen.pdf

Arity.idr

module Arity

import Data.Vect

export
repeat : {n : Nat} -> a -> Vect n a
repeat {n = 0} _ = []
repeat {n = S n} x = x :: repeat {n} x

export
arrTy : {n : Nat} -> Vect (S n) Type -> Type
arrTy {n = 0} [x] = x
arrTy {n = S n} (x :: xs) = x -> arrTy xs

export
zap : Vect n (a -> b) -> Vect n a -> Vect n b
zap [] [] = []
zap (f :: fs) (x :: xs) = f x :: zap fs xs

export
arrTyVec : {n : Nat} -> {m : Nat} -> Vect (S n) Type -> Type
arrTyVec {m} xs = arrTy (zap (repeat (\x => Vect m x)) xs)

export
nvecMap : {m : Nat}
       -> (n : Nat)
       -> (xs : Vect (S n) Type)
       -> arrTy xs
       -> arrTyVec {m} xs
nvecMap n xs f = go {n} xs (repeat f)
  where
    go : {n : Nat} -> {m : Nat} -> (xs : Vect (S n) Type) -> Vect m (arrTy xs) -> arrTyVec {m} xs
    go {n = 0} [_] a = a
    go {n = S n} (x :: xs) f = \a => go xs (zap f a)

export
curryType : {n : Nat} -> {a : Type} -> (Vect n a -> Type) -> Type
curryType {n = 0} b = b []
curryType {n = S n} {a} b = (x : a) -> curryType {n} (\xs => b (x :: xs))

export
curryTerm : {n : Nat}
         -> {a : Type}
         -> {b : Vect n a -> Type}
         -> ((xs : Vect n a) -> b xs)
         -> curryType b
curryTerm {n = 0} f = f []
curryTerm {n = S n} f = \x : a => curryTerm {n} (\xs => f (x :: xs))

export
nmap : (n : Nat)
    -> {m : Nat}
    -> curryType (\as : Vect (S n) Type =>
                   (arrTy as) -> arrTyVec {m} as)
nmap n {m} = curryTerm (nvecMap {m} n)