{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}

-- | Functions on general association (key/value) lists.
module Core.Data.List.Assoc
  ( AssocPair (..)
  , SemAssocPair (..)
  , assocKeys
  , glookup
  , glookupForce
  , assocMember
  , assocUnion
  , assocUnionBy
  , assocInsert
  , assocInsertBy
  , aappend
  ) where

import Data.Semigroup
import Data.List

-- | A generalized key/value pair - can contain more than just the key
-- and value.
class (Eq (Key a)) => AssocPair a where
  type Key a :: *
  type Value a :: *

  -- | The key in the pair.
  getKey :: a -> Key a

  -- | The value in the pair.
  getValue :: a -> Value a

-- | A generalized key/value pair, where 2 pairs with the same key can
-- be appended.
class (AssocPair a) => SemAssocPair a where
  -- | Assumes both values have the same key. Combines them, combining
  -- values with the given function.
  aappend1 :: (Value a -> Value a -> Value a) -> a -> a -> a

-- | All the keys in the association list.
assocKeys :: (AssocPair a) => [a] -> [Key a]
assocKeys = map getKey

-- | Looks up a value in an association list.
glookup :: (AssocPair a) => Key a -> [a] -> Maybe (Value a)
glookup key = fmap getValue . find (hasKey key)

-- | Looks up a value in an association list.
-- Raises an error if the value can't be found.
glookupForce :: (AssocPair a) => Key a -> [a] -> Value a
glookupForce key xs
  = case glookup key xs of
         Nothing -> error "Key not in association list"
         Just val -> val

-- | Is the key in the association list?
assocMember :: (AssocPair a) => Key a -> [a] -> Bool
assocMember = any . hasKey

hasKey :: (AssocPair a) => Key a -> a -> Bool
hasKey key pair = key == getKey pair

-- | Contains all keys from both lists. If 2 values share a key, they're appended.
-- Assumes no pairs share a key in either of the lists before they're combined.
assocUnion :: (SemAssocPair a, Semigroup (Value a)) => [a] -> [a] -> [a]
assocUnion = assocUnionBy (<>)

-- | Contains all keys from both lists. If 2 values share a key, they're appended.
-- Assumes no pairs share a key in either of the lists before they're combined.
assocUnionBy :: (SemAssocPair a)
             => (Value a -> Value a -> Value a)
             -> [a]
             -> [a]
             -> [a]
assocUnionBy = foldr . assocInsertBy

-- | Adds the pair to the list. If the original list already contains
-- its key, the value is appended to the original corresponding value.
assocInsert :: (SemAssocPair a, Monoid (Value a)) => a -> [a] -> [a]
assocInsert = assocInsertBy mappend

-- | Adds the pair to the list. If the original list already contains
-- its key, the value is appended to the original corresponding value.
assocInsertBy :: (SemAssocPair a)
              => (Value a -> Value a -> Value a)
              -> a
              -> [a]
              -> [a]
assocInsertBy _ new [] = [new]
assocInsertBy vapp new (x : xs)
  | getKey new == getKey x = (aappend1 vapp new x) : xs
  | otherwise = x : (assocInsertBy vapp new xs)

-- | Assumes both values have the same key. Combines them.
aappend :: (SemAssocPair a, Monoid (Value a)) => a -> a -> a
aappend = aappend1 mappend