content: {-
    Copyright 2020 The CodeWorld Authors. All rights reserved.
  
    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.
    You may obtain a copy of the License at
  
        http://www.apache.org/licenses/LICENSE-2.0
  
    Unless required by applicable law or agreed to in writing, software
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
    limitations under the License.
  -}
  {-# LANGUAGE TupleSections #-}
  
  -- |
  -- Module       : Data.MultiMap
  -- Copyright    : (c) CodeWorld Authors 2017
  -- License      : Apache-2.0
  --
  -- Maintainer   : Joachim Breitner <mail@joachim-breitner.de>
  --
  -- A simple MultiMap.
  --
  -- This differs from the one in the @multimap@ package by using
  -- 'Data.Sequence.Seq' for efficient insert-at-end and other improved speeds.
  --
  -- Also only supports the operations required by CodeWorld for now.
  module Data.MultiMap
    ( MultiMap,
      empty,
      null,
      insertL,
      insertR,
      toList,
      spanAntitone,
      union,
      keys,
    )
  where
  
  import Data.Bifunctor
  import Data.Coerce
  import qualified Data.Foldable (toList)
  import qualified Data.Map as M
  import qualified Data.Sequence as S
  import Prelude hiding (null)
  
  newtype MultiMap k v
    = MM (M.Map k (S.Seq v))
    deriving (Show, Eq)
  
  empty :: MultiMap k v
  empty = MM M.empty
  
  null :: MultiMap k v -> Bool
  null (MM m) = M.null m
  
  insertL :: Ord k => k -> v -> MultiMap k v -> MultiMap k v
  insertL k v (MM m) = MM (M.alter (Just . maybe (S.singleton v) (v S.<|)) k m)
  
  insertR :: Ord k => k -> v -> MultiMap k v -> MultiMap k v
  insertR k v (MM m) = MM (M.alter (Just . maybe (S.singleton v) (S.|> v)) k m)
  
  toList :: MultiMap k v -> [(k, v)]
  toList (MM m) = [(k, v) | (k, vs) <- M.toList m, v <- Data.Foldable.toList vs]
  
  -- TODO: replace with M.spanAntitone once containers is updated
  mapSpanAntitone :: (k -> Bool) -> M.Map k a -> (M.Map k a, M.Map k a)
  mapSpanAntitone p =
    bimap M.fromDistinctAscList M.fromDistinctAscList
      . span (p . fst)
      . M.toList
  
  spanAntitone :: (k -> Bool) -> MultiMap k v -> (MultiMap k v, MultiMap k v)
  spanAntitone p (MM m) = coerce (mapSpanAntitone p m)
  
  union :: Ord k => MultiMap k v -> MultiMap k v -> MultiMap k v
  union (MM m1) (MM m2) = MM (M.unionWith (S.><) m1 m2)
  
  keys :: MultiMap k v -> [k]
  keys (MM m) = M.keys m
  