SafeFdiv.hs

{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}

-- |
-- Module      :   Grisette.Internal.Core.Data.Class.SafeFdiv
-- Copyright   :   (c) Sirui Lu 2024
-- License     :   BSD-3-Clause (see the LICENSE file)
--
-- Maintainer  :   siruilu@cs.washington.edu
-- Stability   :   Experimental
-- Portability :   GHC only
module Grisette.Internal.Core.Data.Class.SafeFdiv
  ( SafeFdiv (..),
    FdivOr (..),
    fdivOrZero,
    recipOrZero,
  )
where

import Control.Exception (ArithException (RatioZeroDenominator), throw)
import Control.Monad.Error.Class (MonadError (throwError))
import Grisette.Internal.Core.Control.Monad.Class.Union (MonadUnion)
import Grisette.Internal.Core.Data.Class.AsKey (AsKey (AsKey))
import Grisette.Internal.Core.Data.Class.ITEOp (ITEOp (symIte))
import Grisette.Internal.Core.Data.Class.Mergeable (Mergeable)
import Grisette.Internal.Core.Data.Class.SimpleMergeable (mrgIf)
import Grisette.Internal.Core.Data.Class.Solvable (Solvable (con))
import Grisette.Internal.Core.Data.Class.SymEq (SymEq ((.==)))
import Grisette.Internal.Core.Data.Class.TryMerge (TryMerge, mrgSingle, tryMerge)
import Grisette.Internal.SymPrim.AlgReal
  ( AlgReal (AlgExactRational),
    UnsupportedAlgRealOperation (UnsupportedAlgRealOperation),
  )
import Grisette.Internal.SymPrim.SymAlgReal (SymAlgReal)

-- $setup
-- >>> import Grisette.Core
-- >>> import Grisette.SymPrim
-- >>> import Control.Monad.Except
-- >>> import Control.Exception

-- | Safe fractional with default values returned on exception.
class FdivOr a where
  -- | Safe '/' with default values returned on exception.
  --
  -- >>> fdivOr "d" "a" "b" :: SymAlgReal
  -- (ite (= b 0.0) d (fdiv a b))
  fdivOr :: a -> a -> a -> a

  -- | Safe 'recip' with default values returned on exception.
  --
  -- >>> recipOr "d" "a" :: SymAlgReal
  -- (ite (= a 0.0) d (recip a))
  recipOr :: a -> a -> a

-- | Safe '/' with 0 returned on exception.
fdivOrZero :: (FdivOr a, Num a) => a -> a -> a
fdivOrZero l = fdivOr (l - l) l

-- | Safe 'recip' with 0 returned on exception.
recipOrZero :: (FdivOr a, Num a) => a -> a
recipOrZero v = recipOr (v - v) v

-- | Safe fractional division with monadic error handling in multi-path
-- execution. These procedures throw an exception when the denominator is zero.
-- The result should be able to handle errors with `MonadError`.
class (MonadError e m, TryMerge m, Mergeable a) => SafeFdiv e a m where
  -- | Safe fractional division with monadic error handling in multi-path
  -- execution.
  --
  -- >>> safeFdiv "a" "b" :: ExceptT ArithException Union SymAlgReal
  -- ExceptT {If (= b 0.0) (Left Ratio has zero denominator) (Right (fdiv a b))}
  safeFdiv :: a -> a -> m a

  -- | Safe fractional reciprocal with monadic error handling in multi-path
  -- execution.
  --
  -- >>> safeRecip "a" :: ExceptT ArithException Union SymAlgReal
  -- ExceptT {If (= a 0.0) (Left Ratio has zero denominator) (Right (recip a))}
  safeRecip :: a -> m a
  default safeRecip :: (Fractional a) => a -> m a
  safeRecip = safeFdiv (fromRational 1)
  {-# INLINE safeRecip #-}

  {-# MINIMAL safeFdiv #-}

instance FdivOr AlgReal where
  fdivOr d (AlgExactRational l) (AlgExactRational r)
    | r /= 0 = AlgExactRational (l / r)
    | otherwise = d
  fdivOr d l r =
    -- Throw the error because the user should never construct an AlgReal
    -- other than AlgExactRational.
    throw $
      UnsupportedAlgRealOperation "fdivOr" $
        show d <> " and " <> show l <> " and " <> show r
  {-# INLINE fdivOr #-}
  recipOr d (AlgExactRational l)
    | l /= 0 = AlgExactRational (recip l)
    | otherwise = d
  recipOr d l =
    throw $ UnsupportedAlgRealOperation "recipOr" $ show d <> " and " <> show l
  {-# INLINE recipOr #-}

instance
  ( MonadError ArithException m,
    TryMerge m
  ) =>
  SafeFdiv ArithException AlgReal m
  where
  safeFdiv (AlgExactRational l) (AlgExactRational r)
    | r /= 0 =
        pure $ AlgExactRational (l / r)
    | otherwise = tryMerge $ throwError RatioZeroDenominator
  safeFdiv l r =
    -- Throw the error because the user should never construct an AlgReal
    -- other than AlgExactRational.
    throw $
      UnsupportedAlgRealOperation "safeFdiv" $
        show l <> " and " <> show r
  {-# INLINE safeFdiv #-}
  safeRecip (AlgExactRational l)
    | l /= 0 =
        pure $ AlgExactRational (recip l)
    | otherwise = tryMerge $ throwError RatioZeroDenominator
  safeRecip l =
    throw $ UnsupportedAlgRealOperation "safeRecip" $ show l

instance FdivOr SymAlgReal where
  fdivOr d l r = symIte (r .== con 0) d (l / r)
  recipOr d l = symIte (l .== con 0) d (recip l)

instance
  (MonadError ArithException m, MonadUnion m) =>
  SafeFdiv ArithException SymAlgReal m
  where
  safeFdiv l r =
    mrgIf (r .== con 0) (throwError RatioZeroDenominator) (pure $ l / r)
  safeRecip l =
    mrgIf (l .== con 0) (throwError RatioZeroDenominator) (pure $ recip l)

instance (SafeFdiv e a m) => SafeFdiv e (AsKey a) m where
  safeFdiv (AsKey a) (AsKey b) = do
    r <- safeFdiv a b
    mrgSingle $ AsKey r
  safeRecip (AsKey a) = do
    r <- safeRecip a
    mrgSingle $ AsKey r

instance (FdivOr a) => FdivOr (AsKey a) where
  fdivOr (AsKey d) (AsKey a) (AsKey b) = AsKey $ fdivOr d a b
  {-# INLINE fdivOr #-}
  recipOr (AsKey d) (AsKey a) = AsKey $ recipOr d a
  {-# INLINE recipOr #-}