{-----------------------------------------------------------------------------
    reactive-banana

    Implementation of a bag whose elements are ordered by arrival time.
------------------------------------------------------------------------------}
{-# LANGUAGE TupleSections #-}
module Reactive.Banana.Prim.OrderedBag where

import           Data.Functor
import qualified Data.HashMap.Strict as Map
import           Data.Hashable
import           Data.List  hiding (insert)
import           Data.Maybe
import           Data.Ord

{-----------------------------------------------------------------------------
    Ordered Bag
------------------------------------------------------------------------------}
type Position = Integer

data OrderedBag a = OB !(Map.HashMap a Position) !Position

empty :: OrderedBag a
empty :: OrderedBag a
empty = HashMap a Position -> Position -> OrderedBag a
forall a. HashMap a Position -> Position -> OrderedBag a
OB HashMap a Position
forall k v. HashMap k v
Map.empty Position
0

-- | Add an element to an ordered bag after all the others.
-- Does nothing if the element is already in the bag.
insert :: (Eq a, Hashable a) => OrderedBag a -> a -> OrderedBag a
insert :: OrderedBag a -> a -> OrderedBag a
insert (OB HashMap a Position
xs Position
n) a
x = HashMap a Position -> Position -> OrderedBag a
forall a. HashMap a Position -> Position -> OrderedBag a
OB ((Position -> Position -> Position)
-> a -> Position -> HashMap a Position -> HashMap a Position
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
Map.insertWith (\Position
new Position
old -> Position
old) a
x Position
n HashMap a Position
xs) (Position
nPosition -> Position -> Position
forall a. Num a => a -> a -> a
+Position
1)

-- | Add a sequence of elements to an ordered bag.
--
-- The ordering is left-to-right. For example, the head of the sequence
-- comes after all elements in the bag,
-- but before the other elements in the sequence.
inserts :: (Eq a, Hashable a) => OrderedBag a -> [a] -> OrderedBag a
inserts :: OrderedBag a -> [a] -> OrderedBag a
inserts = (OrderedBag a -> a -> OrderedBag a)
-> OrderedBag a -> [a] -> OrderedBag a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' OrderedBag a -> a -> OrderedBag a
forall a. (Eq a, Hashable a) => OrderedBag a -> a -> OrderedBag a
insert

-- | Reorder a list of elements to appear as they were inserted into the bag.
-- Remove any elements from the list that do not appear in the bag.
inOrder :: (Eq a, Hashable a) => [(a,b)] -> OrderedBag a -> [(a,b)]
inOrder :: [(a, b)] -> OrderedBag a -> [(a, b)]
inOrder [(a, b)]
xs (OB HashMap a Position
bag Position
_) = ((Position, (a, b)) -> (a, b)) -> [(Position, (a, b))] -> [(a, b)]
forall a b. (a -> b) -> [a] -> [b]
map (Position, (a, b)) -> (a, b)
forall a b. (a, b) -> b
snd ([(Position, (a, b))] -> [(a, b)])
-> [(Position, (a, b))] -> [(a, b)]
forall a b. (a -> b) -> a -> b
$ ((Position, (a, b)) -> (Position, (a, b)) -> Ordering)
-> [(Position, (a, b))] -> [(Position, (a, b))]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (((Position, (a, b)) -> Position)
-> (Position, (a, b)) -> (Position, (a, b)) -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (Position, (a, b)) -> Position
forall a b. (a, b) -> a
fst) ([(Position, (a, b))] -> [(Position, (a, b))])
-> [(Position, (a, b))] -> [(Position, (a, b))]
forall a b. (a -> b) -> a -> b
$
    ((a, b) -> Maybe (Position, (a, b)))
-> [(a, b)] -> [(Position, (a, b))]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (\(a, b)
x -> (,(a, b)
x) (Position -> (Position, (a, b)))
-> Maybe Position -> Maybe (Position, (a, b))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> HashMap a Position -> Maybe Position
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
Map.lookup ((a, b) -> a
forall a b. (a, b) -> a
fst (a, b)
x) HashMap a Position
bag) [(a, b)]
xs