PartialMonoidOp.hs

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wall -Werror #-}

-- |
-- A partial monoid operation is an associative binary operation with an identity
-- element that is not defined for all pairs of inputs.
module Data.Associative.PartialMonoidOp
  ( -- * Types
    PartialMonoidOp (..),

    -- * Isomorphism
    iPartialMonoidOp,

    -- * Running
    runPartialMonoidOp,
    identityPartialMonoidOp,

    -- * Smart constructors
    pmonoid,

    -- * Laws
    pmonoidLawAssociative,
    pmonoidLawLeftIdentity,
    pmonoidLawRightIdentity,

    -- * Classy optics
    HasPartialMonoidOp (..),
    AsPartialMonoidOp (..),

    -- * Values (via pmonoid)
    pmonoidUnit,
    pmonoidOrdering,
    pmonoidList,
    pmonoidProxy,
    pmonoidMaybe,
    pmonoidDual,
    pmonoidDown,
    pmonoidIdentity,
    pmonoidTuple,
    pmonoidWrappedMonoid,
    pmonoidFunction,
    pmonoidAlt,
    pmonoidAlternative,
    pmonoidLiftF2,
    pmonoidLiftA2,

    -- * Values (via PartialMonoidOp)
    pmonoidMin,
    pmonoidMax,
    pmonoidAll,
    pmonoidAny,
    pmonoidAddition,
    pmonoidMultiplication,
    pmonoidEndo,
    pmonoidAnd,
    pmonoidIor,
    pmonoidXor,
    pmonoidIff,

    -- * Collection values
    pmonoidSetUnion,
    pmonoidIntSetUnion,
    pmonoidHashSetUnion,
    pmonoidMapUnion,
    pmonoidIntMapUnion,
    pmonoidHashMapUnion,
  )
where

import Control.Applicative (Alternative (..))
import Control.Lens
  ( Iso,
    Lens',
    Prism',
    iso,
    lens,
    review,
  )
import Data.Associative.MonoidOp (MonoidOp (..))
import Data.Associative.PartialSemigroupOp
  ( HasPartialSemigroupOpT (..),
    PartialSemigroupOp',
    iPartialSemigroupOp,
    psemigroupDown,
    psemigroupDual,
    psemigroupFunction,
    psemigroupIdentity,
    psemigroupLiftA2,
    psemigroupMaybe,
    psemigroupTuple,
    psemigroupWrappedMonoid,
    runPartialSemigroupOp,
    psemigroupSemigroup,
    total,
  )
import Data.Associative.SemigroupOp (SemigroupOp')
import Data.Bits (Bits, FiniteBits, complement, xor, zeroBits, (.&.), (.|.))
import Data.Functor.Alt (Alt (..))
import Data.Functor.Identity (Identity (..))
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HashMap
import Data.HashSet (HashSet)
import qualified Data.HashSet as HashSet
import Data.Hashable (Hashable)
import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import Data.IntSet (IntSet)
import qualified Data.IntSet as IntSet
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Ord (Down (..))
import Data.Proxy (Proxy (..))
import Data.Semigroup (Dual (..), WrappedMonoid (..))
import Data.Set (Set)
import qualified Data.Set as Set
import GHC.Generics (Generic)

-- $setup
-- >>> import Data.Associative.SemigroupOp (SemigroupOp', op, semigroupList)
-- >>> import Data.Associative.MonoidOp (MonoidOp(..), monoidList)
-- >>> import Data.Associative.PartialSemigroupOp (PartialSemigroupOpT(..), total)
-- >>> import Control.Lens (view, review)
-- >>> import Data.Functor.Identity (Identity(..))
-- >>> import Data.Ord (Down(..))
-- >>> import Data.Proxy (Proxy(..))
-- >>> import Data.Semigroup (Dual(..), WrappedMonoid(..))
-- >>> import Data.Word (Word8)
-- >>> import qualified Data.Set as Set
-- >>> import qualified Data.IntSet as IntSet
-- >>> import qualified Data.HashSet as HashSet
-- >>> import qualified Data.Map as Map
-- >>> import qualified Data.IntMap as IntMap
-- >>> import qualified Data.HashMap.Strict as HashMap
-- >>> import Data.List (sort)
-- >>> let add = PartialMonoidOp (total (+)) 0 :: PartialMonoidOp Int
-- >>> let run = runPartialMonoidOp

-- | A partial monoid operation: an associative binary operation with an identity
-- element, not necessarily defined for all pairs of inputs.
--
-- >>> run add 3 4
-- Just 7
-- >>> identityPartialMonoidOp add
-- 0
data PartialMonoidOp a = PartialMonoidOp (PartialSemigroupOp' a) a
  deriving (Generic)

-- | Iso between 'PartialMonoidOp' and a @(partial-binary-operation, identity)@ pair.
--
-- >>> let (f, e) = view iPartialMonoidOp add
-- >>> f 3 4
-- Just 7
-- >>> e
-- 0
iPartialMonoidOp :: Iso (PartialMonoidOp a) (PartialMonoidOp b) (a -> a -> Maybe a, a) (b -> b -> Maybe b, b)
iPartialMonoidOp =
  iso
    (\(PartialMonoidOp s e) -> (runPartialSemigroupOp s, e))
    (\(f, e) -> PartialMonoidOp (review iPartialSemigroupOp f) e)
{-# INLINE iPartialMonoidOp #-}

-- | Extract the partial binary operation and run it.
--
-- >>> runPartialMonoidOp add 3 4
-- Just 7
runPartialMonoidOp :: PartialMonoidOp a -> a -> a -> Maybe a
runPartialMonoidOp (PartialMonoidOp s _) = runPartialSemigroupOp s
{-# INLINE runPartialMonoidOp #-}

-- | Extract the identity element.
--
-- >>> identityPartialMonoidOp add
-- 0
-- >>> identityPartialMonoidOp pmonoidList
-- []
identityPartialMonoidOp :: PartialMonoidOp a -> a
identityPartialMonoidOp (PartialMonoidOp _ e) = e
{-# INLINE identityPartialMonoidOp #-}

-- | Build a 'PartialMonoidOp' from any 'Monoid' instance.
--
-- >>> run (pmonoid :: PartialMonoidOp [Int]) [1,2] [3,4]
-- Just [1,2,3,4]
-- >>> identityPartialMonoidOp (pmonoid :: PartialMonoidOp [Int])
-- []
pmonoid :: (Monoid a) => PartialMonoidOp a
pmonoid = PartialMonoidOp psemigroupSemigroup mempty
{-# INLINE pmonoid #-}

-- | Classy lens giving access to the underlying 'PartialSemigroupOpT'.
--
-- >>> import Data.Associative.PartialSemigroupOp (runPartialSemigroupOp)
-- >>> runPartialSemigroupOp (view partialSemigroupOpT add) 3 4
-- Just 7
instance HasPartialSemigroupOpT (PartialMonoidOp a) Identity a a where
  partialSemigroupOpT = lens (\(PartialMonoidOp s _) -> s) (\(PartialMonoidOp _ e) s' -> PartialMonoidOp s' e)

{- HLINT ignore "Monoid law, left identity" -}
{- HLINT ignore "Monoid law, right identity" -}

----
-- Law-checking functions
----

-- | Associativity of the partial monoid operation.
--
-- Left- and right-association of three values must agree:
-- both 'Nothing' or both the same 'Just'.
--
-- >>> pmonoidLawAssociative add 1 2 3
-- True
pmonoidLawAssociative :: (Eq a) => PartialMonoidOp a -> a -> a -> a -> Bool
pmonoidLawAssociative m x y z =
  let f = runPartialMonoidOp m
      lhs = case f x y of
        Nothing -> Nothing
        Just xy -> f xy z
      rhs = case f y z of
        Nothing -> Nothing
        Just yz -> f x yz
   in lhs == rhs

-- | Left identity: @f e a == Just a@
--
-- >>> pmonoidLawLeftIdentity add 42
-- True
-- >>> pmonoidLawLeftIdentity pmonoidList [1,2,3 :: Int]
-- True
pmonoidLawLeftIdentity :: (Eq a) => PartialMonoidOp a -> a -> Bool
pmonoidLawLeftIdentity m a =
  runPartialMonoidOp m (identityPartialMonoidOp m) a == Just a

-- | Right identity: @f a e == Just a@
--
-- >>> pmonoidLawRightIdentity add 42
-- True
-- >>> pmonoidLawRightIdentity pmonoidList [1,2,3 :: Int]
-- True
pmonoidLawRightIdentity :: (Eq a) => PartialMonoidOp a -> a -> Bool
pmonoidLawRightIdentity m a =
  runPartialMonoidOp m a (identityPartialMonoidOp m) == Just a

-- | Classy lens for types that contain a 'PartialMonoidOp'.
--
-- >>> run (view partialMonoidOp add) 3 4
-- Just 7
class HasPartialMonoidOp c a | c -> a where
  partialMonoidOp :: Lens' c (PartialMonoidOp a)

instance HasPartialMonoidOp (PartialMonoidOp a) a where
  partialMonoidOp = id

-- | Classy prism for types that can be constructed from a 'PartialMonoidOp'.
--
-- >>> run (review _PartialMonoidOp add) 3 4
-- Just 7
class AsPartialMonoidOp c a | c -> a where
  _PartialMonoidOp :: Prism' c (PartialMonoidOp a)

instance AsPartialMonoidOp (PartialMonoidOp a) a where
  _PartialMonoidOp = id

----
-- PartialMonoidOp values via pmonoid
----

-- | >>> run pmonoidUnit () ()
-- Just ()
-- >>> identityPartialMonoidOp pmonoidUnit
-- ()
pmonoidUnit :: PartialMonoidOp ()
pmonoidUnit = pmonoid

-- | Lexicographic composition of orderings.
--
-- >>> run pmonoidOrdering LT GT
-- Just LT
-- >>> run pmonoidOrdering EQ GT
-- Just GT
-- >>> identityPartialMonoidOp pmonoidOrdering
-- EQ
pmonoidOrdering :: PartialMonoidOp Ordering
pmonoidOrdering = pmonoid

-- | List concatenation.
--
-- >>> run pmonoidList [1,2] [3,4 :: Int]
-- Just [1,2,3,4]
-- >>> identityPartialMonoidOp pmonoidList
-- []
pmonoidList :: PartialMonoidOp [a]
pmonoidList = pmonoid

-- | >>> run pmonoidProxy Proxy (Proxy :: Proxy Int)
-- Just Proxy
-- >>> identityPartialMonoidOp pmonoidProxy
-- Proxy
pmonoidProxy :: PartialMonoidOp (Proxy a)
pmonoidProxy = pmonoid

-- | 'Nothing' is identity; 'Just' values are combined.
--
-- >>> run (pmonoidMaybe semigroupList) (Just [1]) (Just [2 :: Int])
-- Just (Just [1,2])
-- >>> run (pmonoidMaybe semigroupList) Nothing (Just [2 :: Int])
-- Just (Just [2])
-- >>> identityPartialMonoidOp (pmonoidMaybe semigroupList)
-- Nothing
pmonoidMaybe :: SemigroupOp' a -> PartialMonoidOp (Maybe a)
pmonoidMaybe s = PartialMonoidOp (psemigroupMaybe s) Nothing

-- | Reverses the inner monoid.
--
-- >>> run (pmonoidDual monoidList) (Dual [1]) (Dual [2 :: Int])
-- Just (Dual {getDual = [2,1]})
-- >>> identityPartialMonoidOp (pmonoidDual (monoidList :: MonoidOp [Int]))
-- Dual {getDual = []}
pmonoidDual :: MonoidOp a -> PartialMonoidOp (Dual a)
pmonoidDual (MonoidOp s e) = PartialMonoidOp (psemigroupDual s) (Dual e)

-- | Delegates through 'Down'.
--
-- >>> run (pmonoidDown monoidList) (Down [1]) (Down [2 :: Int])
-- Just (Down [1,2])
pmonoidDown :: MonoidOp a -> PartialMonoidOp (Down a)
pmonoidDown (MonoidOp s e) = PartialMonoidOp (psemigroupDown s) (Down e)

-- | Delegates through 'Identity'.
--
-- >>> run (pmonoidIdentity monoidList) (Identity [1]) (Identity [2 :: Int])
-- Just (Identity [1,2])
pmonoidIdentity :: MonoidOp a -> PartialMonoidOp (Identity a)
pmonoidIdentity (MonoidOp s e) = PartialMonoidOp (psemigroupIdentity s) (Identity e)

-- | Pairwise combination.
--
-- >>> run (pmonoidTuple monoidList monoidList) ([1 :: Int], [10]) ([2], [20 :: Int])
-- Just ([1,2],[10,20])
-- >>> identityPartialMonoidOp (pmonoidTuple monoidList monoidList :: PartialMonoidOp ([Int], [Int]))
-- ([],[])
pmonoidTuple :: MonoidOp a -> MonoidOp b -> PartialMonoidOp (a, b)
pmonoidTuple (MonoidOp sa ea) (MonoidOp sb eb) = PartialMonoidOp (psemigroupTuple sa sb) (ea, eb)

-- | Uses the underlying monoid operation.
--
-- >>> run (pmonoidWrappedMonoid monoidList) (WrapMonoid [1]) (WrapMonoid [2 :: Int])
-- Just (WrapMonoid {unwrapMonoid = [1,2]})
pmonoidWrappedMonoid :: MonoidOp a -> PartialMonoidOp (WrappedMonoid a)
pmonoidWrappedMonoid (MonoidOp s e) = PartialMonoidOp (psemigroupWrappedMonoid s) (WrapMonoid e)

-- | Pointwise combination.
--
-- >>> fmap ($ "x") (run (pmonoidFunction monoidList) (++ "a") ((++ "b") :: String -> String))
-- Just "xaxb"
pmonoidFunction :: MonoidOp b -> PartialMonoidOp (a -> b)
pmonoidFunction (MonoidOp s e) = PartialMonoidOp (psemigroupFunction s) (const e)

-- | First-success on 'Maybe' via 'Alt'.
--
-- >>> run pmonoidAlt (Just 1) (Just 2 :: Maybe Int)
-- Just (Just 1)
-- >>> run pmonoidAlt Nothing (Just 2 :: Maybe Int)
-- Just (Just 2)
-- >>> identityPartialMonoidOp pmonoidAlt
-- Nothing
pmonoidAlt :: PartialMonoidOp (Maybe a)
pmonoidAlt = PartialMonoidOp (total (<!>)) Nothing

-- | First-success on 'Maybe' via 'Alternative'.
--
-- >>> run pmonoidAlternative (Just 1) (Just 2 :: Maybe Int)
-- Just (Just 1)
-- >>> run pmonoidAlternative Nothing (Just 2 :: Maybe Int)
-- Just (Just 2)
-- >>> identityPartialMonoidOp pmonoidAlternative
-- Nothing
pmonoidAlternative :: PartialMonoidOp (Maybe a)
pmonoidAlternative = PartialMonoidOp (total (<|>)) Nothing

-- | Lift a monoid operation through an 'Data.Functor.Apply.Apply' functor via 'Data.Functor.Apply.liftF2'.
-- Requires 'Applicative' for 'pure' to construct the identity element.
--
-- >>> run (pmonoidLiftF2 monoidList) (Just [1,2]) (Just [3,4 :: Int])
-- Just (Just [1,2,3,4])
-- >>> identityPartialMonoidOp (pmonoidLiftF2 monoidList :: PartialMonoidOp (Maybe [Int]))
-- Just []
pmonoidLiftF2 :: (Applicative f) => MonoidOp a -> PartialMonoidOp (f a)
pmonoidLiftF2 (MonoidOp s e) = PartialMonoidOp (psemigroupLiftA2 s) (pure e)

-- | Lift a monoid operation through an 'Applicative' functor via 'Control.Applicative.liftA2'.
--
-- >>> run (pmonoidLiftA2 monoidList) (Just [1,2]) (Just [3,4 :: Int])
-- Just (Just [1,2,3,4])
-- >>> identityPartialMonoidOp (pmonoidLiftA2 monoidList :: PartialMonoidOp (Maybe [Int]))
-- Just []
pmonoidLiftA2 :: (Applicative f) => MonoidOp a -> PartialMonoidOp (f a)
pmonoidLiftA2 (MonoidOp s e) = PartialMonoidOp (psemigroupLiftA2 s) (pure e)

----
-- PartialMonoidOp values via PartialMonoidOp constructor
----

-- | Takes the minimum ('Min'). Requires 'Bounded' for 'maxBound' identity.
--
-- >>> run pmonoidMin (3 :: Int) 4
-- Just 3
-- >>> identityPartialMonoidOp pmonoidMin == (maxBound :: Int)
-- True
pmonoidMin :: (Ord a, Bounded a) => PartialMonoidOp a
pmonoidMin = PartialMonoidOp (total min) maxBound

-- | Takes the maximum ('Max'). Requires 'Bounded' for 'minBound' identity.
--
-- >>> run pmonoidMax (3 :: Int) 4
-- Just 4
-- >>> identityPartialMonoidOp pmonoidMax == (minBound :: Int)
-- True
pmonoidMax :: (Ord a, Bounded a) => PartialMonoidOp a
pmonoidMax = PartialMonoidOp (total max) minBound

-- | Logical conjunction ('All'). Identity is 'True'.
--
-- >>> run pmonoidAll True True
-- Just True
-- >>> run pmonoidAll True False
-- Just False
-- >>> identityPartialMonoidOp pmonoidAll
-- True
pmonoidAll :: PartialMonoidOp Bool
pmonoidAll = PartialMonoidOp (total (&&)) True

-- | Logical disjunction ('Any'). Identity is 'False'.
--
-- >>> run pmonoidAny False False
-- Just False
-- >>> run pmonoidAny False True
-- Just True
-- >>> identityPartialMonoidOp pmonoidAny
-- False
pmonoidAny :: PartialMonoidOp Bool
pmonoidAny = PartialMonoidOp (total (||)) False

-- | Addition ('Sum'). Identity is 0.
--
-- >>> run pmonoidAddition (3 :: Int) 4
-- Just 7
-- >>> identityPartialMonoidOp pmonoidAddition
-- 0
pmonoidAddition :: (Num a) => PartialMonoidOp a
pmonoidAddition = PartialMonoidOp (total (+)) 0

-- | Multiplication ('Product'). Identity is 1.
--
-- >>> run pmonoidMultiplication (3 :: Int) 4
-- Just 12
-- >>> identityPartialMonoidOp pmonoidMultiplication
-- 1
pmonoidMultiplication :: (Num a) => PartialMonoidOp a
pmonoidMultiplication = PartialMonoidOp (total (*)) 1

-- | Function composition ('Endo'). Identity is 'id'.
--
-- >>> fmap ($ 3) (run pmonoidEndo (+1) ((*10) :: Int -> Int))
-- Just 31
-- >>> identityPartialMonoidOp pmonoidEndo 42
-- 42
pmonoidEndo :: PartialMonoidOp (a -> a)
pmonoidEndo = PartialMonoidOp (total (.)) id

-- | Bitwise AND. Identity is all ones ('complement' 'zeroBits').
--
-- >>> run pmonoidAnd (0xFF :: Word8) 0x0F
-- Just 15
-- >>> identityPartialMonoidOp pmonoidAnd == (0xFF :: Word8)
-- True
pmonoidAnd :: (Bits a) => PartialMonoidOp a
pmonoidAnd = PartialMonoidOp (total (.&.)) (complement zeroBits)

-- | Bitwise inclusive OR. Identity is 'zeroBits'.
--
-- >>> run pmonoidIor (0xF0 :: Word8) 0x0F
-- Just 255
-- >>> identityPartialMonoidOp pmonoidIor == (0 :: Word8)
-- True
pmonoidIor :: (Bits a) => PartialMonoidOp a
pmonoidIor = PartialMonoidOp (total (.|.)) zeroBits

-- | Bitwise exclusive OR. Identity is 'zeroBits'.
--
-- >>> run pmonoidXor (0xFF :: Word8) 0x0F
-- Just 240
-- >>> identityPartialMonoidOp pmonoidXor == (0 :: Word8)
-- True
pmonoidXor :: (Bits a) => PartialMonoidOp a
pmonoidXor = PartialMonoidOp (total xor) zeroBits

-- | Bitwise equivalence / XNOR. Identity is all ones ('complement' 'zeroBits').
--
-- >>> run pmonoidIff (0xFF :: Word8) 0x0F
-- Just 15
-- >>> identityPartialMonoidOp pmonoidIff == (0xFF :: Word8)
-- True
pmonoidIff :: (FiniteBits a) => PartialMonoidOp a
pmonoidIff = PartialMonoidOp (total (\a b -> complement (xor a b))) (complement zeroBits)

----
-- Collection values
----

-- | Set union. Identity is 'Set.empty'.
--
-- >>> run pmonoidSetUnion (Set.fromList [1,2]) (Set.fromList [2,3 :: Int])
-- Just (fromList [1,2,3])
-- >>> identityPartialMonoidOp pmonoidSetUnion == (Set.empty :: Set Int)
-- True
pmonoidSetUnion :: (Ord a) => PartialMonoidOp (Set a)
pmonoidSetUnion = PartialMonoidOp (total Set.union) Set.empty

-- | IntSet union. Identity is 'IntSet.empty'.
--
-- >>> run pmonoidIntSetUnion (IntSet.fromList [1,2]) (IntSet.fromList [2,3])
-- Just (fromList [1,2,3])
-- >>> identityPartialMonoidOp pmonoidIntSetUnion == IntSet.empty
-- True
pmonoidIntSetUnion :: PartialMonoidOp IntSet
pmonoidIntSetUnion = PartialMonoidOp (total IntSet.union) IntSet.empty

-- | HashSet union. Identity is 'HashSet.empty'.
--
-- >>> fmap sort (fmap HashSet.toList (run pmonoidHashSetUnion (HashSet.fromList [1,2]) (HashSet.fromList [2,3 :: Int])))
-- Just [1,2,3]
pmonoidHashSetUnion :: (Eq a, Hashable a) => PartialMonoidOp (HashSet a)
pmonoidHashSetUnion = PartialMonoidOp (total HashSet.union) HashSet.empty

-- | Map union (left-biased on overlapping keys). Identity is 'Map.empty'.
--
-- >>> run pmonoidMapUnion (Map.fromList [(1 :: Int,'a'),(2,'b')]) (Map.fromList [(2,'x'),(3,'c')])
-- Just (fromList [(1,'a'),(2,'b'),(3,'c')])
pmonoidMapUnion :: (Ord k) => PartialMonoidOp (Map k v)
pmonoidMapUnion = PartialMonoidOp (total Map.union) Map.empty

-- | IntMap union (left-biased on overlapping keys). Identity is 'IntMap.empty'.
--
-- >>> run pmonoidIntMapUnion (IntMap.fromList [(1,'a'),(2,'b')]) (IntMap.fromList [(2,'x'),(3,'c')])
-- Just (fromList [(1,'a'),(2,'b'),(3,'c')])
pmonoidIntMapUnion :: PartialMonoidOp (IntMap v)
pmonoidIntMapUnion = PartialMonoidOp (total IntMap.union) IntMap.empty

-- | HashMap union (left-biased on overlapping keys). Identity is 'HashMap.empty'.
--
-- >>> fmap sort (fmap HashMap.toList (run pmonoidHashMapUnion (HashMap.fromList [(1 :: Int,'a'),(2,'b')]) (HashMap.fromList [(2,'x'),(3,'c')])))
-- Just [(1,'a'),(2,'b'),(3,'c')]
pmonoidHashMapUnion :: (Eq k, Hashable k) => PartialMonoidOp (HashMap k v)
pmonoidHashMapUnion = PartialMonoidOp (total HashMap.union) HashMap.empty