Ix.hs

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

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

module Data.Foldable.Ix where

import Data.Foldable
import GHC.Base hiding (foldr, foldl')
import GHC.Num
import GHC.Real
import GHC.List (take, drop, reverse)

data OneInTwoBang a b = B12 !a !b deriving Eq

data ThreeInFourBang a b = B34 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 = (\(B12 n1 _) -> if n1 == (-1) then Nothing else Just n1) . foldl' f v $ js
  where v = B12 (-1) 0
        f (B12 n m) !y
         | n >= 0 = B12 n (m + 1)
         | y == x = B12 m (m + 1)
         | otherwise = B12 n (m + 1)
{-# SPECIALIZE findIdx1 :: (Eq a, Foldable t) => a -> t a -> Maybe Int #-}

findIdx1' ::  (Eq a, Foldable t) => a -> t a -> Maybe Int 
findIdx1' = findIdx1
       
{-| 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 = (\(B12 ys _) -> ys) . foldr f v $ js
  where v = B12 [] (length js - 1)
        f y (B12 xs m)
         | y == x = B12 (m:xs) (m - 1)
         | otherwise = B12 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 = (\(B12 ys _) -> reverse ys) . foldl' f v $ js
  where v = B12 [] 0
        f (B12 xs m) !y
         | y == x = B12 (m:xs) (m + 1)
         | otherwise = B12 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 l 0 []
        f (B34 l i xs) x
         | i >= idx && i <= idx + l - 1 = B34 l (i+1) (x:xs)
         | otherwise = B34 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 #-}