Graded semigroups/monoids/groups - two versions

graded1.hs

{-# language FlexibleInstances #-}
{-# language DefaultSignatures #-}
{-# language RankNTypes #-}
{-# language MultiParamTypeClasses #-}
{-# language QuantifiedConstraints #-}
module Data.Group.Graded where

import Data.Functor.WithIndex
import Data.Group
import Data.Map (Map)
import qualified Data.Map as M

class GradedSemigroup i f where
  iappend :: Semigroup g => i -> g -> f g -> f g
  default iappend :: (Eq i, FunctorWithIndex i f, Semigroup g) => i -> g -> f g -> f g
  iappend i h = imap go where
    go j g
      | i == j = g <> h
      | otherwise = g
  {-# inline iappend #-}

class GradedSemigroup i f => GradedMonoid i f where
  imempty :: Monoid g => i -> f g -> f g
  default imempty :: (Monoid g, Eq i) => i -> f g -> f g
  imempty _ fg = fg
  {-# inline imempty #-}

class GradedMonoid i f => GradedGroup i f where
  iinvert :: Group g => i -> f g -> f g
  default iinvert :: (Eq i, FunctorWithIndex i f, Group g) => i -> f g -> f g
  iinvert i = imap go where
    go j g
      | i == j = invert g
      | otherwise = g
  {-# inline iinvert #-}

instance Ord k => GradedSemigroup k (Map k) where
  iappend = M.insertWith (<>)
  {-# inline iappend #-}
  
instance Ord k => GradedMonoid k (Map k)

instance Ord k => GradedGroup k (Map k) where
  iinvert = M.adjust invert
  {-# inline iinvert #-}

instance GradedSemigroup Int []
instance GradedMonoid Int []
instance GradedGroup Int []

-- -------------------------------------------------------------------- --
-- Dense indexed groups (no mempty's)

newtype DenseGradedMap k v = DenseGradedMap { runDenseGradedMap :: Map k v }

fromMap :: (Monoid v, Ord k, Eq v) => Map k v -> DenseGradedMap k v
fromMap = DenseGradedMap . M.foldrWithKey go mempty
  where
    go _ IdentityElem acc = acc
    go k g acc = M.insert k g acc
{-# inline fromMap #-}

fromMapUnsafe :: Map k v -> DenseGradedMap k v
fromMapUnsafe = DenseGradedMap
{-# inline fromMapUnsafe #-}

instance (Ord k, Monoid v, Eq v) => Semigroup (DenseGradedMap k v) where
  -- this needs to be improved. Should be hedge-union and key removal on mempty for /O(n + m)/
  DenseGradedMap m <> DenseGradedMap n = DenseGradedMap $ M.foldrWithKey go mempty (m <> n)
    where
      go _ IdentityElem acc = acc
      go k g acc = M.insert k g acc


instance (Ord k, Monoid v, Eq v) => Monoid (DenseGradedMap k v) where
  mempty = DenseGradedMap mempty

instance (Ord k, Group v, Eq v) => Group (DenseGradedMap k v) where
  invert (DenseGradedMap m) = DenseGradedMap (M.map invert m)

instance Ord k => GradedSemigroup k (DenseGradedMap k) where
  -- this should prune. gotta thonks about it.
  iappend i g (DenseGradedMap m) = DenseGradedMap $ M.insertWith (<>) i g m

instance Ord k => GradedMonoid k (DenseGradedMap k) where
  imempty _ fg = fg

instance Ord k => GradedGroup k (DenseGradedMap k) where
  iinvert i (DenseGradedMap m) = DenseGradedMap $ M.adjust invert i m