{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeFamilies #-}
module Fresnel.Ixed
( -- * Indexable collections
  Ixed(..)
  -- * Construction
, ixSet
, ixMap
, ixList
) where

import           Control.Monad (guard)
import qualified Data.HashMap.Internal as HashMap
import qualified Data.HashSet as HashSet
import           Data.Hashable (Hashable)
import qualified Data.IntMap as IntMap
import qualified Data.IntSet as IntSet
import qualified Data.List.NonEmpty as NonEmpty
import qualified Data.Map as Map
import qualified Data.Set as Set
import           Fresnel.List.NonEmpty (head_, tail_)
import           Fresnel.Optional (Optional', optional')

-- Indexable collections

class Ixed c where
  type Index c
  type IxValue c

  ix :: Index c -> Optional' c (IxValue c)

instance Ixed IntSet.IntSet where
  type Index IntSet.IntSet = IntSet.Key
  type IxValue IntSet.IntSet = ()

  ix :: Index IntSet -> Optional' IntSet (IxValue IntSet)
ix = (Index IntSet -> IntSet -> Bool)
-> (Index IntSet -> IntSet -> IntSet)
-> Index IntSet
-> Optional' IntSet ()
forall c.
(Index c -> c -> Bool)
-> (Index c -> c -> c) -> Index c -> Optional' c ()
ixSet Key -> IntSet -> Bool
Index IntSet -> IntSet -> Bool
IntSet.member Key -> IntSet -> IntSet
Index IntSet -> IntSet -> IntSet
IntSet.insert

instance Ixed (IntMap.IntMap v) where
  type Index (IntMap.IntMap v) = IntMap.Key
  type IxValue (IntMap.IntMap v) = v

  ix :: Index (IntMap v) -> Optional' (IntMap v) (IxValue (IntMap v))
ix = (Index (IntMap v) -> IntMap v -> Maybe (IxValue (IntMap v)))
-> (Index (IntMap v) -> IxValue (IntMap v) -> IntMap v -> IntMap v)
-> Index (IntMap v)
-> Optional' (IntMap v) (IxValue (IntMap v))
forall c.
(Index c -> c -> Maybe (IxValue c))
-> (Index c -> IxValue c -> c -> c)
-> Index c
-> Optional' c (IxValue c)
ixMap Index (IntMap v) -> IntMap v -> Maybe (IxValue (IntMap v))
forall a. Key -> IntMap a -> Maybe a
IntMap.lookup Index (IntMap v) -> IxValue (IntMap v) -> IntMap v -> IntMap v
forall a. Key -> a -> IntMap a -> IntMap a
IntMap.insert

instance Ord k => Ixed (Set.Set k) where
  type Index (Set.Set k) = k
  type IxValue (Set.Set k) = ()

  ix :: Index (Set k) -> Optional' (Set k) (IxValue (Set k))
ix = (Index (Set k) -> Set k -> Bool)
-> (Index (Set k) -> Set k -> Set k)
-> Index (Set k)
-> Optional' (Set k) ()
forall c.
(Index c -> c -> Bool)
-> (Index c -> c -> c) -> Index c -> Optional' c ()
ixSet Index (Set k) -> Set k -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Index (Set k) -> Set k -> Set k
forall a. Ord a => a -> Set a -> Set a
Set.insert

instance Ord k => Ixed (Map.Map k v) where
  type Index (Map.Map k v) = k
  type IxValue (Map.Map k v) = v

  ix :: Index (Map k v) -> Optional' (Map k v) (IxValue (Map k v))
ix = (Index (Map k v) -> Map k v -> Maybe (IxValue (Map k v)))
-> (Index (Map k v) -> IxValue (Map k v) -> Map k v -> Map k v)
-> Index (Map k v)
-> Optional' (Map k v) (IxValue (Map k v))
forall c.
(Index c -> c -> Maybe (IxValue c))
-> (Index c -> IxValue c -> c -> c)
-> Index c
-> Optional' c (IxValue c)
ixMap Index (Map k v) -> Map k v -> Maybe (IxValue (Map k v))
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Index (Map k v) -> IxValue (Map k v) -> Map k v -> Map k v
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert

instance (Eq k, Hashable k) => Ixed (HashSet.HashSet k) where
  type Index (HashSet.HashSet k) = k
  type IxValue (HashSet.HashSet k) = ()

  ix :: Index (HashSet k) -> Optional' (HashSet k) (IxValue (HashSet k))
ix = (Index (HashSet k) -> HashSet k -> Bool)
-> (Index (HashSet k) -> HashSet k -> HashSet k)
-> Index (HashSet k)
-> Optional' (HashSet k) ()
forall c.
(Index c -> c -> Bool)
-> (Index c -> c -> c) -> Index c -> Optional' c ()
ixSet Index (HashSet k) -> HashSet k -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
HashSet.member Index (HashSet k) -> HashSet k -> HashSet k
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
HashSet.insert

instance (Eq k, Hashable k) => Ixed (HashMap.HashMap k v) where
  type Index (HashMap.HashMap k v) = k
  type IxValue (HashMap.HashMap k v) = v

  ix :: Index (HashMap k v)
-> Optional' (HashMap k v) (IxValue (HashMap k v))
ix = (Index (HashMap k v)
 -> HashMap k v -> Maybe (IxValue (HashMap k v)))
-> (Index (HashMap k v)
    -> IxValue (HashMap k v) -> HashMap k v -> HashMap k v)
-> Index (HashMap k v)
-> Optional' (HashMap k v) (IxValue (HashMap k v))
forall c.
(Index c -> c -> Maybe (IxValue c))
-> (Index c -> IxValue c -> c -> c)
-> Index c
-> Optional' c (IxValue c)
ixMap Index (HashMap k v) -> HashMap k v -> Maybe (IxValue (HashMap k v))
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HashMap.lookup Index (HashMap k v)
-> IxValue (HashMap k v) -> HashMap k v -> HashMap k v
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HashMap.insert

instance Ixed (Maybe a) where
  type Index (Maybe a) = ()
  type IxValue (Maybe a) = a

  ix :: Index (Maybe a) -> Optional' (Maybe a) (IxValue (Maybe a))
ix Index (Maybe a)
_ = (Maybe a -> Maybe a)
-> (Maybe a -> a -> Maybe a) -> Optional (Maybe a) (Maybe a) a a
forall s a b. (s -> Maybe a) -> (s -> b -> s) -> Optional s s a b
optional' Maybe a -> Maybe a
forall a. a -> a
id (\ Maybe a
_ a
a -> a -> Maybe a
forall a. a -> Maybe a
Just a
a)

instance Ixed [v] where
  type Index [v] = Int
  type IxValue [v] = v

  ix :: Index [v] -> Optional' [v] (IxValue [v])
ix Index [v]
k = Key -> Optional' [v] v
forall a. Key -> Optional' [a] a
ixList Key
Index [v]
k

instance Ixed (NonEmpty.NonEmpty v) where
  type Index (NonEmpty.NonEmpty v) = Int
  type IxValue (NonEmpty.NonEmpty v) = v

  ix :: Index (NonEmpty v) -> Optional' (NonEmpty v) (IxValue (NonEmpty v))
ix Index (NonEmpty v)
k
    | Key
Index (NonEmpty v)
k Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
<= Key
0    = p (IxValue (NonEmpty v)) (IxValue (NonEmpty v))
-> p (NonEmpty v) (NonEmpty v)
forall a. Lens' (NonEmpty a) a
head_
    | Bool
otherwise = Optic p (NonEmpty v) (NonEmpty v) [v] [v]
forall a. Lens' (NonEmpty a) [a]
tail_Optic p (NonEmpty v) (NonEmpty v) [v] [v]
-> (p v v -> p [v] [v]) -> p v v -> p (NonEmpty v) (NonEmpty v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.Key -> Optional' [v] v
forall a. Key -> Optional' [a] a
ixList (Key
Index (NonEmpty v)
k Key -> Key -> Key
forall a. Num a => a -> a -> a
- Key
1)


-- Construction

ixSet :: (Index c -> c -> Bool) -> (Index c -> c -> c) -> Index c -> Optional' c ()
ixSet :: (Index c -> c -> Bool)
-> (Index c -> c -> c) -> Index c -> Optional' c ()
ixSet Index c -> c -> Bool
member Index c -> c -> c
insert Index c
k = (c -> Maybe ()) -> (c -> () -> c) -> Optional' c ()
forall s a b. (s -> Maybe a) -> (s -> b -> s) -> Optional s s a b
optional' (Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Bool -> Maybe ()) -> (c -> Bool) -> c -> Maybe ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Index c -> c -> Bool
member Index c
k) (c -> () -> c
forall a b. a -> b -> a
const (c -> () -> c) -> (c -> c) -> c -> () -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Index c -> c -> c
insert Index c
k)

ixMap :: (Index c -> c -> Maybe (IxValue c)) -> (Index c -> IxValue c -> c -> c) -> Index c -> Optional' c (IxValue c)
ixMap :: (Index c -> c -> Maybe (IxValue c))
-> (Index c -> IxValue c -> c -> c)
-> Index c
-> Optional' c (IxValue c)
ixMap Index c -> c -> Maybe (IxValue c)
lookup Index c -> IxValue c -> c -> c
insert Index c
k = (c -> Maybe (IxValue c))
-> (c -> IxValue c -> c) -> Optional' c (IxValue c)
forall s a b. (s -> Maybe a) -> (s -> b -> s) -> Optional s s a b
optional' (Index c -> c -> Maybe (IxValue c)
lookup Index c
k) ((IxValue c -> c -> c) -> c -> IxValue c -> c
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Index c -> IxValue c -> c -> c
insert Index c
k))

ixList :: Int -> Optional' [a] a
ixList :: Key -> Optional' [a] a
ixList Key
i = ([a] -> Maybe a) -> ([a] -> a -> [a]) -> Optional' [a] a
forall s a b. (s -> Maybe a) -> (s -> b -> s) -> Optional s s a b
optional' (Key -> [a] -> Maybe a
forall t a. (Ord t, Num t) => t -> [a] -> Maybe a
get Key
i) (Key -> [a] -> a -> [a]
forall t t. (Ord t, Num t) => t -> [t] -> t -> [t]
set Key
i)
  where
  get :: t -> [a] -> Maybe a
get t
i [a]
as = case [a]
as of
    []   -> Maybe a
forall a. Maybe a
Nothing
    a
a:[a]
as -> if t
i t -> t -> Bool
forall a. Ord a => a -> a -> Bool
<= t
0 then a -> Maybe a
forall a. a -> Maybe a
Just a
a else t -> [a] -> Maybe a
get (t
i t -> t -> t
forall a. Num a => a -> a -> a
- t
1) [a]
as
  set :: t -> [t] -> t -> [t]
set t
i [t]
as t
a' = case [t]
as of
    []   -> [t]
as
    t
a:[t]
as -> if t
i t -> t -> Bool
forall a. Ord a => a -> a -> Bool
<= t
0 then t
a't -> [t] -> [t]
forall a. a -> [a] -> [a]
:[t]
as else t
a t -> [t] -> [t]
forall a. a -> [a] -> [a]
: t -> [t] -> t -> [t]
set (t
i t -> t -> t
forall a. Num a => a -> a -> a
- t
1) [t]
as t
a'