{-# 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 OneInTwoBang a b -> OneInTwoBang a b -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall a b.
(Eq a, Eq b) =>
OneInTwoBang a b -> OneInTwoBang a b -> Bool
/= :: OneInTwoBang a b -> OneInTwoBang a b -> Bool
$c/= :: forall a b.
(Eq a, Eq b) =>
OneInTwoBang a b -> OneInTwoBang a b -> Bool
== :: OneInTwoBang a b -> OneInTwoBang a b -> Bool
$c== :: forall a b.
(Eq a, Eq b) =>
OneInTwoBang a b -> OneInTwoBang a b -> Bool
Eq

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

findIdx1' ::  (Eq a, Foldable t) => a -> t a -> Maybe Int 
findIdx1' :: forall a (t :: * -> *). (Eq a, Foldable t) => a -> t a -> Maybe Int
findIdx1' = forall a (t :: * -> *) b.
(Eq a, Foldable t, Integral b) =>
a -> t a -> Maybe b
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 :: forall a (t :: * -> *). (Eq a, Foldable t) => a -> t a -> [Int]
findIdxs a
x t a
js = (\(B12 [Int]
ys Int
_) -> [Int]
ys) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr forall {b}. Num b => a -> OneInTwoBang [b] b -> OneInTwoBang [b] b
f forall {a}. OneInTwoBang [a] Int
v forall a b. (a -> b) -> a -> b
$ t a
js
  where v :: OneInTwoBang [a] Int
v = forall a b. a -> b -> OneInTwoBang a b
B12 [] (forall (t :: * -> *) a. Foldable t => t a -> Int
length t a
js forall a. Num a => a -> a -> a
- Int
1)
        f :: a -> OneInTwoBang [b] b -> OneInTwoBang [b] b
f a
y (B12 [b]
xs b
m)
         | a
y forall a. Eq a => a -> a -> Bool
== a
x = forall a b. a -> b -> OneInTwoBang a b
B12 (b
mforall a. a -> [a] -> [a]
:[b]
xs) (b
m forall a. Num a => a -> a -> a
- b
1)
         | Bool
otherwise = forall a b. a -> b -> OneInTwoBang a b
B12 [b]
xs (b
m forall a. Num a => a -> a -> a
- b
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 :: forall a (t :: * -> *). (Eq a, Foldable t) => a -> t a -> [Int]
findIdxsL1 a
x t a
js = (\(B12 [Int]
ys Int
_) -> forall a. [a] -> [a]
reverse [Int]
ys) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall {b}. Num b => OneInTwoBang [b] b -> a -> OneInTwoBang [b] b
f forall {a}. OneInTwoBang [a] Int
v forall a b. (a -> b) -> a -> b
$ t a
js
  where v :: OneInTwoBang [a] Int
v = forall a b. a -> b -> OneInTwoBang a b
B12 [] Int
0
        f :: OneInTwoBang [b] b -> a -> OneInTwoBang [b] b
f (B12 [b]
xs b
m) !a
y
         | a
y forall a. Eq a => a -> a -> Bool
== a
x = forall a b. a -> b -> OneInTwoBang a b
B12 (b
mforall a. a -> [a] -> [a]
:[b]
xs) (b
m forall a. Num a => a -> a -> a
+ b
1)
         | Bool
otherwise = forall a b. a -> b -> OneInTwoBang a b
B12 [b]
xs (b
m forall a. Num a => a -> a -> a
+ b
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 :: forall a (t :: * -> *).
(Eq a, Foldable t) =>
Int -> Int -> t a -> [a]
sliceToList Int
idx Int
l t a
js = (\(B34 Int
_ Int
_ [a]
ys) -> forall a. [a] -> [a]
reverse [a]
ys) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall {a}. ThreeInFourBang a Int -> a -> ThreeInFourBang a Int
f forall {a}. ThreeInFourBang a Int
v forall a b. (a -> b) -> a -> b
$ t a
js
  where v :: ThreeInFourBang a Int
v = forall a b. b -> b -> [a] -> ThreeInFourBang a b
B34 Int
l Int
0 []
        f :: ThreeInFourBang a Int -> a -> ThreeInFourBang a Int
f (B34 Int
l Int
i [a]
xs) a
x
         | Int
i forall a. Ord a => a -> a -> Bool
>= Int
idx Bool -> Bool -> Bool
&& Int
i forall a. Ord a => a -> a -> Bool
<= Int
idx forall a. Num a => a -> a -> a
+ Int
l forall a. Num a => a -> a -> a
- Int
1 = forall a b. b -> b -> [a] -> ThreeInFourBang a b
B34 Int
l (Int
iforall a. Num a => a -> a -> a
+Int
1) (a
xforall a. a -> [a] -> [a]
:[a]
xs)
         | Bool
otherwise = forall a b. b -> b -> [a] -> ThreeInFourBang a b
B34 Int
l (Int
iforall a. Num a => a -> a -> a
+Int
1) [a]
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 :: forall a. Eq a => Int -> Int -> [a] -> [a]
s2L Int
idx Int
l = forall a. Int -> [a] -> [a]
drop Int
idx forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Int -> [a] -> [a]
take (Int
idx forall a. Num a => a -> a -> a
+ Int
l)
{-# INLINABLE s2L #-}