module TraversableUtil where

import Data.FunctorShape
import Data.Traversable (mapAccumL, mapAccumM)
import Data.Foldable (toList)

indices :: Traversable f => f x -> f Int
indices :: forall (f :: * -> *) x. Traversable f => f x -> f Int
indices = (Int, f Int) -> f Int
forall a b. (a, b) -> b
snd ((Int, f Int) -> f Int) -> (f x -> (Int, f Int)) -> f x -> f Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> x -> (Int, Int)) -> Int -> f x -> (Int, f Int)
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL (\Int
n x
_ -> (Int -> Int
forall a. Enum a => a -> a
succ Int
n, Int
n)) Int
0

-- >>> zipMatch [1,2,3] "abc"
-- Just [(1,'a'),(2,'b'),(3,'c')]
-- >>> zipMatch [1,2] "abc"
-- Nothing
zipMatch :: (Eq (t Ignored), Traversable t) => t a -> t b -> Maybe (t (a,b))
zipMatch :: forall (t :: * -> *) a b.
(Eq (t Ignored), Traversable t) =>
t a -> t b -> Maybe (t (a, b))
zipMatch t a
ta t b
tb
  | t a -> Shape t
forall (f :: * -> *) a. f a -> Shape f
Shape t a
ta Shape t -> Shape t -> Bool
forall a. Eq a => a -> a -> Bool
== t b -> Shape t
forall (f :: * -> *) a. f a -> Shape f
Shape t b
tb = ([b], t (a, b)) -> t (a, b)
forall a b. (a, b) -> b
snd (([b], t (a, b)) -> t (a, b))
-> Maybe ([b], t (a, b)) -> Maybe (t (a, b))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([b] -> a -> Maybe ([b], (a, b)))
-> [b] -> t a -> Maybe ([b], t (a, b))
forall (m :: * -> *) (t :: * -> *) s a b.
(Monad m, Traversable t) =>
(s -> a -> m (s, b)) -> s -> t a -> m (s, t b)
mapAccumM [b] -> a -> Maybe ([b], (a, b))
forall {b} {a}. [b] -> a -> Maybe ([b], (a, b))
step (t b -> [b]
forall a. t a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList t b
tb) t a
ta
  | Bool
otherwise = Maybe (t (a, b))
forall a. Maybe a
Nothing
  where
    step :: [b] -> a -> Maybe ([b], (a, b))
step [] a
_ = Maybe ([b], (a, b))
forall a. Maybe a
Nothing
    step (b
b:[b]
bs) a
a = ([b], (a, b)) -> Maybe ([b], (a, b))
forall a. a -> Maybe a
Just ([b]
bs, (a
a,b
b))