Vector.hs

module Composite.Ix.Vector (groupBy, groupBy', groupBy'') where

import           Composite.Record
import           Control.Lens
import qualified Data.Map         as M
import           Data.Proxy
import           Data.Vector      as V (Vector, singleton, toList)
import           Data.Vinyl       hiding (rlens, rlens')
import           Prelude

-- | Group a vector of records by a subtype of the record type.
--
-- @since 0.0.1.0
groupBy :: forall ys xs f.
           (ys ⊆ xs)
        => Ord (Rec f ys)
        => Vector (Rec f xs)
        -> M.Map (Rec f ys) (Vector (Rec f xs))
groupBy xs = M.fromListWith (<>) $ V.toList $ fmap (\x -> (rcast x, V.singleton x)) xs

-- | Group a vector of records by a single field.
--
-- @since 0.0.1.0
groupBy' :: forall k s y xs.
            k ~ (s :-> y)
         => k ∈ xs
         => Ord y
         => Vector (Record xs)
         -> M.Map y (Vector (Record xs))
groupBy' xs = M.fromListWith (<>) $ V.toList $ fmap (\x -> (view (rlens (Proxy @(s :-> y))) x, V.singleton x)) xs

-- | Group a vector of records by a single field (HKD).
--
-- @since 0.0.1.0
groupBy'' :: forall k s y xs f.
             k ~ (s :-> y)
          => k ∈ xs
          => Functor f
          => Ord (f y)
          => Vector (Rec f xs)
          -> M.Map (f y) (Vector (Rec f xs))
groupBy'' xs = M.fromListWith (<>) $ V.toList $ fmap (\x -> (view (rlens' (Proxy @(s :-> y))) x, V.singleton x)) xs