{-# LANGUAGE  ScopedTypeVariables  #-}
--------------------------------------------------------------------------------
-- |
-- Module      :  HGeometry.Indexed
-- Copyright   :  (C) Frank Staals
-- License     :  see the LICENSE file
-- Maintainer  :  Frank Staals
--
-- Things that have an index.
--
--------------------------------------------------------------------------------
module HGeometry.Indexed
  ( HasIndex(..)
  , Index
  , WithIndex(..), theValue
  , labelWithIndex, labelWith, labelWith'
  ) where


import Control.Lens(Lens, Field1, lens, _1)
import Control.Monad.State.Strict

--------------------------------------------------------------------------------

-- | A type of Indices
type Index = Int

-- | Types that have an index.
class HasIndex a where
  -- | Get the index of a given 'a'
  sosIndex :: a -> Index

-- | Associate a value of type 'a' with an Index.
data WithIndex a = WithIndex {-# UNPACK #-} !Index a
               deriving (Int -> WithIndex a -> ShowS
[WithIndex a] -> ShowS
WithIndex a -> String
(Int -> WithIndex a -> ShowS)
-> (WithIndex a -> String)
-> ([WithIndex a] -> ShowS)
-> Show (WithIndex a)
forall a. Show a => Int -> WithIndex a -> ShowS
forall a. Show a => [WithIndex a] -> ShowS
forall a. Show a => WithIndex a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> WithIndex a -> ShowS
showsPrec :: Int -> WithIndex a -> ShowS
$cshow :: forall a. Show a => WithIndex a -> String
show :: WithIndex a -> String
$cshowList :: forall a. Show a => [WithIndex a] -> ShowS
showList :: [WithIndex a] -> ShowS
Show)

-- | Lens to manipulate the value of the 'WithIndex'
theValue :: Lens (WithIndex a) (WithIndex b) a b
theValue :: forall a b (f :: * -> *).
Functor f =>
(a -> f b) -> WithIndex a -> f (WithIndex b)
theValue = (WithIndex a -> a)
-> (WithIndex a -> b -> WithIndex b)
-> Lens (WithIndex a) (WithIndex b) a b
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens (\(WithIndex Int
_ a
x) -> a
x) (\(WithIndex Int
i a
_) b
y -> Int -> b -> WithIndex b
forall a. Int -> a -> WithIndex a
WithIndex Int
i b
y)

instance Field1 (WithIndex a) (WithIndex b) a b where
  _1 :: Lens (WithIndex a) (WithIndex b) a b
_1 = (a -> f b) -> WithIndex a -> f (WithIndex b)
forall a b (f :: * -> *).
Functor f =>
(a -> f b) -> WithIndex a -> f (WithIndex b)
theValue

instance HasIndex (WithIndex a) where
  sosIndex :: WithIndex a -> Int
sosIndex (WithIndex Int
i a
_) = Int
i
  {-# INLINE sosIndex #-}

-- instance Eq a => Eq (WithIndex a) where
--   (WithIndex i x) == (WithIndex j y) = x == y && i == j

-- instance Ord a => Ord (WithIndex a) where
--   (WithIndex i x) `compare` (WithIndex j y) = x `compare` y <> Down i `compare` Down j

--------------------------------------------------------------------------------


--------------------------------------------------------------------------------

-- | Label each element with its index.
labelWithIndex :: Traversable  t => t a -> t (WithIndex a)
labelWithIndex :: forall (t :: * -> *) a. Traversable t => t a -> t (WithIndex a)
labelWithIndex = (Int -> a -> WithIndex a) -> t a -> t (WithIndex a)
forall (t :: * -> *) a b.
Traversable t =>
(Int -> a -> b) -> t a -> t b
labelWith Int -> a -> WithIndex a
forall a. Int -> a -> WithIndex a
WithIndex

-- | Label each element with its index using the given labelling
-- function.
labelWith   :: Traversable t => (Index -> a -> b) -> t a -> t b
labelWith :: forall (t :: * -> *) a b.
Traversable t =>
(Int -> a -> b) -> t a -> t b
labelWith Int -> a -> b
f = (t b, Int) -> t b
forall a b. (a, b) -> a
fst ((t b, Int) -> t b) -> (t a -> (t b, Int)) -> t a -> t b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> a -> b) -> t a -> (t b, Int)
forall (t :: * -> *) a b.
Traversable t =>
(Int -> a -> b) -> t a -> (t b, Int)
labelWith' Int -> a -> b
f

-- | Label each element with its index using the given labelling
-- function. Returns the new collection as well as its size.
labelWith'           :: forall t a b. Traversable t
                     => (Index -> a -> b) -> t a
                     -> (t b, Int)
labelWith' :: forall (t :: * -> *) a b.
Traversable t =>
(Int -> a -> b) -> t a -> (t b, Int)
labelWith' Int -> a -> b
withIndex = (State Int (t b) -> Int -> (t b, Int))
-> Int -> State Int (t b) -> (t b, Int)
forall a b c. (a -> b -> c) -> b -> a -> c
flip State Int (t b) -> Int -> (t b, Int)
forall s a. State s a -> s -> (a, s)
runState Int
0 (State Int (t b) -> (t b, Int))
-> (t a -> State Int (t b)) -> t a -> (t b, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> StateT Int Identity b) -> t a -> State Int (t b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> t a -> f (t b)
traverse a -> StateT Int Identity b
lbl
  where
    lbl   :: a -> State Int b
    lbl :: a -> StateT Int Identity b
lbl a
x = do i <- StateT Int Identity Int
forall s (m :: * -> *). MonadState s m => m s
get
               put $ i+1
               pure (withIndex i x)