{-# LANGUAGE RankNTypes #-}
module Data.List.Extras (wix) where

import Control.Lens
import Data.Word

-- | A variant of 'ix' that takes a 'Word64' instead of an 'Int'.
wix :: Word64 -> Traversal' [a] a
wix :: forall a. Word64 -> Traversal' [a] a
wix Word64
k a -> f a
f [a]
xs0 = [a] -> Word64 -> f [a]
forall {a}. (Eq a, Num a) => [a] -> a -> f [a]
go [a]
xs0 Word64
k where
    go :: [a] -> a -> f [a]
go [] a
_     = [a] -> f [a]
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    go (a
a:[a]
as) a
0 = a -> f a
f a
a f a -> (a -> [a]) -> f [a]
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> (a -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
as)
    go (a
a:[a]
as) a
i = (a
aa -> [a] -> [a]
forall a. a -> [a] -> [a]
:) ([a] -> [a]) -> f [a] -> f [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([a] -> a -> f [a]
go [a]
as (a -> f [a]) -> a -> f [a]
forall a b. (a -> b) -> a -> b
$! a
i a -> a -> a
forall a. Num a => a -> a -> a
- a
1)