{-# LANGUAGE CPP, BangPatterns #-}
{-# OPTIONS_HADDOCK show-extensions #-}
{-# OPTIONS_GHC -funbox-strict-fields #-}

-- |
-- Module      :  Data.Foldable.Ix
-- Copyright   :  (c) OleksandrZhabenko 2020-2021
-- License     :  MIT
-- Stability   :  Experimental
-- Maintainer  :  olexandr543@yahoo.com
--
--

module Data.Foldable.Ix where

import Data.Foldable

data TwoInThreeBang a b = B23 a !Int !Int deriving Eq

data TwoInThreeBang2 a = B23L a ![Int] !Int

data ThreeInFourBang a b = B34 b b !b ![a] deriving Eq

{-| Function to find out the \'index\' (as the reperesentative of the 'Integral' class) of the first element in the 'Foldable' structure (from the left with indices starting from 0), which equals to the first argument. Returns 'Nothing' if there are no such elements.
-}
findIdx1 :: (Eq a, Foldable t, Integral b) => a -> t a -> Maybe b
findIdx1 x js = (\(_, n1, _) -> if n1 == (-1) then Nothing else Just n1) . foldl' f v $ js
  where v = (x, (-1), 0)
        f (t, n, m) y
         | n >= 0 = (t, n, m + 1)
         | y == t = (t, m, m + 1)
         | otherwise = (t, n, m + 1)

{-| A variant of the 'findIdx1' where the resulting 'Maybe' b is 'Maybe' 'Int'. Possibly can be more optimized.
-}
findIdx1' :: (Eq a, Foldable t) => a -> t a -> Maybe Int
findIdx1' x js = (\(B23 _ n1 _) -> if n1 == (-1) then Nothing else Just n1) . foldl' f v $ js
  where v = B23 x (-1) 0
        f (B23 t n m) y
         | n >= 0 = B23 t n (m + 1)
         | y == t = B23 t m (m + 1)
         | otherwise = B23 t n (m + 1)
         
{-| Function to find out the \'indices\' of the elements in the 'Foldable' structure (from the left with indices starting from 0) that equal to the first argument. Returns empty list if there are no such elements. Uses two passes
through the structure.
-}
findIdxs :: (Eq a, Foldable t) => a -> t a -> [Int]
findIdxs x js = (\(_,ys,_) -> ys) . foldr f v $ js
  where v = (x,[],length js - 1)
        f y (t,xs,m)
         | y == t = (t,m:xs,m - 1)
         | otherwise = (t,xs,m - 1)

{-| Function to find out the \'indices\' of the elements in the 'Foldable' structure (from the left with indices starting from 0) that equal to the first argument. Returns empty list if there are no such elements. Uses just one
pass through the structure and additional 'reverse' operation on the resulting list with 'foldl''.
-}
findIdxsL1 :: (Eq a, Foldable t) => a -> t a -> [Int]
findIdxsL1 x js = (\(B23L _ ys _) -> reverse ys) . foldl' f v $ js
  where v = B23L x [] 0
        f (B23L t xs m) y
         | y == t = B23L t (m:xs) (m + 1)
         | otherwise = B23L t xs (m + 1)

{-| Inspired by the Data.Vector.slice function from the @vector@ package. Takes a \'slice\' for the 'Foldable' structure converting it to the list. The first argument is the \'index\' of the element in the structure starting from 0 from the left. The second one is the length of the slice.
-}
sliceToList :: (Eq a, Foldable t) => Int -> Int -> t a -> [a]
sliceToList idx l js = (\(B34 _ _ _ ys) -> reverse ys) . foldl' f v $ js
  where v = B34 idx l 0 []
        f (B34 idx l i xs) x
         | i >= idx && i <= idx + l - 1 = B34 idx l (i+1) (x:xs)
         | otherwise = B34 idx l (i+1) xs
{-# SPECIALIZE sliceToList :: (Eq a) => Int -> Int -> [a] -> [a] #-}
{-# NOINLINE[2] sliceToList #-}

{-# RULES "sliceToList/lists" sliceToList = s2L #-}
s2L :: (Eq a) => Int -> Int -> [a] -> [a]
s2L idx l = drop idx . take (idx + l)
{-# INLINABLE s2L #-}