content: {-# LANGUAGE CPP #-}
  {-# LANGUAGE ConstraintKinds #-}
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE DerivingVia #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE QuantifiedConstraints #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE TypeOperators #-}
  {-# LANGUAGE UndecidableInstances #-}
  
  -- |
  -- Module      :   Grisette.Internal.Core.Data.Class.SafeSymShift
  -- Copyright   :   (c) Sirui Lu 2021-2023
  -- License     :   BSD-3-Clause (see the LICENSE file)
  --
  -- Maintainer  :   siruilu@cs.washington.edu
  -- Stability   :   Experimental
  -- Portability :   GHC only
  module Grisette.Internal.Core.Data.Class.SafeSymShift
    ( SafeSymShift (..),
    )
  where
  
  import Control.Exception (ArithException (Overflow))
  import Control.Monad.Error.Class (MonadError (throwError))
  import Data.Bits (Bits (shiftL, shiftR), FiniteBits (finiteBitSize))
  import Data.Int (Int16, Int32, Int64, Int8)
  import Data.Word (Word16, Word32, Word64, Word8)
  import GHC.TypeLits (KnownNat, type (<=))
  import Grisette.Internal.Core.Control.Monad.Class.Union (MonadUnion)
  import Grisette.Internal.Core.Data.Class.LogicalOp
    ( LogicalOp ((.&&), (.||)),
    )
  import Grisette.Internal.Core.Data.Class.Mergeable (Mergeable)
  import Grisette.Internal.Core.Data.Class.SimpleMergeable
    ( mrgIf,
    )
  import Grisette.Internal.Core.Data.Class.SymOrd
    ( SymOrd ((.<), (.>=)),
    )
  import Grisette.Internal.Core.Data.Class.TryMerge
    ( TryMerge,
      mrgSingle,
      tryMerge,
    )
  import Grisette.Internal.SymPrim.BV (IntN, WordN)
  import Grisette.Internal.SymPrim.Prim.Term
    ( PEvalShiftTerm
        ( pevalShiftLeftTerm,
          pevalShiftRightTerm
        ),
    )
  import Grisette.Internal.SymPrim.SymBV (SymIntN (SymIntN), SymWordN (SymWordN))
  
  -- | Safe version for `shiftL` or `shiftR`.
  --
  -- The `safeSymShiftL` and `safeSymShiftR` are defined for all non-negative
  -- shift amounts.
  --
  -- * Shifting by negative shift amounts is an error.
  -- * The result is defined to be 0 when shifting left by more than or equal to
  -- the  bit size of the number.
  -- * The result is defined to be 0 when shifting right by more than or equal to
  -- the bit size of the number and the number is unsigned or signed non-negative.
  -- * The result is defined to be -1 when shifting right by more than or equal to
  -- the bit size of the number and the number is signed negative.
  --
  -- The `safeSymStrictShiftL` and `safeSymStrictShiftR` are defined for all
  -- non-negative shift amounts that is less than the bit size. Shifting by more
  -- than or equal to the bit size is an error, otherwise they are the same as the
  -- non-strict versions.
  class (MonadError e m, TryMerge m, Mergeable a) => SafeSymShift e a m where
    safeSymShiftL :: a -> a -> m a
    safeSymShiftR :: a -> a -> m a
    safeSymStrictShiftL :: a -> a -> m a
    safeSymStrictShiftR :: a -> a -> m a
  
  -- | This function handles the case when the shift amount is out the range of
  -- `Int` correctly.
  safeSymShiftLConcreteNum ::
    (MonadError ArithException m, TryMerge m, Integral a, FiniteBits a, Mergeable a) =>
    Bool ->
    a ->
    a ->
    m a
  safeSymShiftLConcreteNum _ _ s | s < 0 = tryMerge $ throwError Overflow
  safeSymShiftLConcreteNum allowLargeShiftAmount a s
    | (fromIntegral s :: Integer) >= fromIntegral (finiteBitSize a) =
        if allowLargeShiftAmount
          then mrgSingle 0
          else tryMerge $ throwError Overflow
  safeSymShiftLConcreteNum _ a s = mrgSingle $ shiftL a (fromIntegral s)
  
  -- | This function handles the case when the shift amount is out the range of
  -- `Int` correctly.
  safeSymShiftRConcreteNum ::
    ( MonadError ArithException m,
      TryMerge m,
      Integral a,
      FiniteBits a,
      Mergeable a
    ) =>
    Bool ->
    a ->
    a ->
    m a
  safeSymShiftRConcreteNum _ _ s | s < 0 = tryMerge $ throwError Overflow
  safeSymShiftRConcreteNum allowLargeShiftAmount a s
    | (fromIntegral s :: Integer) >= fromIntegral (finiteBitSize a) =
        if allowLargeShiftAmount
          then mrgSingle 0
          else tryMerge $ throwError Overflow
  safeSymShiftRConcreteNum _ a s = mrgSingle $ shiftR a (fromIntegral s)
  
  #define SAFE_SYM_SHIFT_CONCRETE(T) \
    instance (MonadError ArithException m, TryMerge m) => \
      SafeSymShift ArithException T m where \
      safeSymShiftL = safeSymShiftLConcreteNum True; \
      safeSymShiftR = safeSymShiftRConcreteNum True; \
      safeSymStrictShiftL = safeSymShiftLConcreteNum False; \
      safeSymStrictShiftR = safeSymShiftRConcreteNum False
  
  #if 1
  SAFE_SYM_SHIFT_CONCRETE(Word8)
  SAFE_SYM_SHIFT_CONCRETE(Word16)
  SAFE_SYM_SHIFT_CONCRETE(Word32)
  SAFE_SYM_SHIFT_CONCRETE(Word64)
  SAFE_SYM_SHIFT_CONCRETE(Word)
  SAFE_SYM_SHIFT_CONCRETE(Int8)
  SAFE_SYM_SHIFT_CONCRETE(Int16)
  SAFE_SYM_SHIFT_CONCRETE(Int32)
  SAFE_SYM_SHIFT_CONCRETE(Int64)
  SAFE_SYM_SHIFT_CONCRETE(Int)
  #endif
  
  instance
    (MonadError ArithException m, TryMerge m, KnownNat n, 1 <= n) =>
    SafeSymShift ArithException (WordN n) m
    where
    safeSymShiftL = safeSymShiftLConcreteNum True
    safeSymShiftR = safeSymShiftRConcreteNum True
    safeSymStrictShiftL = safeSymShiftLConcreteNum False
    safeSymStrictShiftR = safeSymShiftRConcreteNum False
  
  instance
    (MonadError ArithException m, TryMerge m, KnownNat n, 1 <= n) =>
    SafeSymShift ArithException (IntN n) m
    where
    safeSymShiftL = safeSymShiftLConcreteNum True
    safeSymShiftR = safeSymShiftRConcreteNum True
    safeSymStrictShiftL = safeSymShiftLConcreteNum False
    safeSymStrictShiftR = safeSymShiftRConcreteNum False
  
  instance
    (MonadError ArithException m, MonadUnion m, KnownNat n, 1 <= n) =>
    SafeSymShift ArithException (SymWordN n) m
    where
    safeSymShiftL (SymWordN a) (SymWordN s) =
      return $ SymWordN $ pevalShiftLeftTerm a s
    safeSymShiftR (SymWordN a) (SymWordN s) =
      return $ SymWordN $ pevalShiftRightTerm a s
    safeSymStrictShiftL a@(SymWordN ta) s@(SymWordN ts) =
      mrgIf
        (s .>= fromIntegral (finiteBitSize a))
        (throwError Overflow)
        (return $ SymWordN $ pevalShiftLeftTerm ta ts)
    safeSymStrictShiftR a@(SymWordN ta) s@(SymWordN ts) =
      mrgIf
        (s .>= fromIntegral (finiteBitSize a))
        (throwError Overflow)
        (return $ SymWordN $ pevalShiftRightTerm ta ts)
  
  instance
    (MonadError ArithException m, MonadUnion m, KnownNat n, 1 <= n) =>
    SafeSymShift ArithException (SymIntN n) m
    where
    safeSymShiftL (SymIntN a) ss@(SymIntN s) =
      mrgIf
        (ss .< 0)
        (throwError Overflow)
        (return $ SymIntN $ pevalShiftLeftTerm a s)
    safeSymShiftR (SymIntN a) ss@(SymIntN s) =
      mrgIf
        (ss .< 0)
        (throwError Overflow)
        (return $ SymIntN $ pevalShiftRightTerm a s)
    safeSymStrictShiftL a@(SymIntN ta) s@(SymIntN ts) =
      mrgIf
        (s .< 0 .|| (bs .>= 0 .&& s .>= bs))
        (throwError Overflow)
        (return $ SymIntN $ pevalShiftLeftTerm ta ts)
      where
        bs = fromIntegral (finiteBitSize a)
    safeSymStrictShiftR a@(SymIntN ta) s@(SymIntN ts) =
      mrgIf
        (s .< 0 .|| (bs .>= 0 .&& s .>= bs))
        (throwError Overflow)
        (return $ SymIntN $ pevalShiftRightTerm ta ts)
      where
        bs = fromIntegral (finiteBitSize a)
  