MapF

MapF.hs

{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE FlexibleInstances #-}
module MyMap where

import qualified Data.Map as M
import qualified Data.Map.Merge.Lazy as M
import Control.Applicative
import Control.Comonad
import Data.Functor.Identity

data MapF f k a = MapF
    { def :: f a
    , mapping :: M.Map k (f a)
    } deriving (Show, Eq, Functor, Foldable, Traversable)

instance (Ord k, Applicative f) => Applicative (MapF f k) where
  pure a = MapF (pure a) M.empty
  MapF defF mapF <*> MapF defA mapA =
      MapF (defF <*> defA)
           (M.merge (M.mapMissing (const (<*> defA)))
                    (M.mapMissing (const (defF <*>)))
                    (M.zipWithMatched (const (<*>)))
                    mapF
                    mapA
           )

instance (Ord k, Alternative f) => Alternative (MapF f k) where
  empty = MapF empty M.empty
  MapF defA mapA <|> MapF defB mapB = MapF (defA <|> defB) (M.unionWith (<|>) mapA mapB)


instance (Ord k, Comonad f) => Comonad (MapF f k) where
  duplicate :: forall a. MapF f k a -> MapF f k (MapF f k a)
  duplicate w@(MapF d m) = MapF (extend (\d' -> (MapF d' m)) d) (M.mapWithKey go m)
    where
      go :: k -> f a -> f (MapF f k a)
      go k fa = extend (\fa' -> MapF fa' (M.delete k m)) fa
  extract = extract . def

instance (Ord k, Semigroup (f a)) => Semigroup (MapF f k a)  where
  MapF defA mapA <> MapF defB mapB = MapF (defA <> defB) (M.unionWith (<>) mapA mapB)

instance (Ord k, Monoid (f a)) => Monoid (MapF f k a)  where
  mempty = MapF mempty M.empty