content: {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  
  module Test.Massiv.Core.SchedulerSpec (spec) where
  
  import Control.Exception.Base (ArithException (DivideByZero))
  import Data.Massiv.Array as A
  import Test.Massiv.Core
  import Prelude as P
  
  -- | Ensure proper exception handling.
  prop_CatchDivideByZero :: ArrIx D Ix2 Int -> [Int] -> Property
  prop_CatchDivideByZero (ArrIx arr ix) caps =
    assertDeepException
      (== DivideByZero)
      ( A.sum $
          A.imap
            ( \ix' x ->
                if ix == ix'
                  then x `div` 0
                  else x
            )
            (setComp (ParOn caps) arr)
      )
  
  -- | Ensure proper exception handling in nested parallel computation
  prop_CatchNested :: ArrIx D Ix1 (ArrIx D Ix1 Int) -> [Int] -> Property
  prop_CatchNested (ArrIx arr ix) caps =
    assertDeepException
      (== DivideByZero)
      ( computeAs U $
          A.map A.sum $
            A.imap
              ( \ix' (ArrIx iarr ixi) ->
                  if ix == ix'
                    then
                      A.imap
                        ( \ixi' e ->
                            if ixi == ixi'
                              then e `div` 0
                              else e
                        )
                        iarr
                    else iarr
              )
              (setComp (ParOn caps) arr)
      )
  
  spec :: Spec
  spec =
    describe "Scheduler - Exceptions" $ do
      it "CatchDivideByZero" $ property prop_CatchDivideByZero
      it "CatchNested" $ property prop_CatchNested
  