{-# LANGUAGE BlockArguments        #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternGuards         #-}
{-# LANGUAGE PatternSynonyms       #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TemplateHaskell       #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE ViewPatterns          #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
-- |
-- Module      : Data.Array.Accelerate.Data.Maybe
-- Copyright   : [2018..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--
-- @since 1.2.0.0
--

module Data.Array.Accelerate.Data.Maybe (

  Maybe(..), pattern Nothing_, pattern Just_,
  maybe, isJust, isNothing, fromMaybe, fromJust, justs,

) where

import Data.Array.Accelerate.AST.Idx
import Data.Array.Accelerate.Language
import Data.Array.Accelerate.Lift
import Data.Array.Accelerate.Pattern.Maybe
import Data.Array.Accelerate.Prelude
import Data.Array.Accelerate.Smart
import Data.Array.Accelerate.Sugar.Array                            ( Array, Vector )
import Data.Array.Accelerate.Sugar.Elt
import Data.Array.Accelerate.Sugar.Shape                            ( Shape, Slice, (:.) )
import Data.Array.Accelerate.Type

import Data.Array.Accelerate.Classes.Eq
import Data.Array.Accelerate.Classes.Ord

import Data.Array.Accelerate.Data.Functor
import Data.Array.Accelerate.Data.Monoid
import Data.Array.Accelerate.Data.Semigroup

import Data.Maybe                                                   ( Maybe(..) )
import Prelude                                                      ( ($), (.) )


-- | Returns 'True' if the argument is 'Nothing'
--
isNothing :: Elt a => Exp (Maybe a) -> Exp Bool
isNothing :: Exp (Maybe a) -> Exp Bool
isNothing = Exp Bool -> Exp Bool
not (Exp Bool -> Exp Bool)
-> (Exp (Maybe a) -> Exp Bool) -> Exp (Maybe a) -> Exp Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp (Maybe a) -> Exp Bool
forall a. Elt a => Exp (Maybe a) -> Exp Bool
isJust

-- | Returns 'True' if the argument is of the form @Just _@
--
isJust :: Elt a => Exp (Maybe a) -> Exp Bool
isJust :: Exp (Maybe a) -> Exp Bool
isJust (Exp SmartExp (EltR (Maybe a))
x) = SmartExp (EltR Bool) -> Exp Bool
forall t. SmartExp (EltR t) -> Exp t
Exp (SmartExp (EltR Bool) -> Exp Bool)
-> SmartExp (EltR Bool) -> Exp Bool
forall a b. (a -> b) -> a -> b
$ PreSmartExp SmartAcc SmartExp (TAG, ()) -> SmartExp (TAG, ())
forall t. PreSmartExp SmartAcc SmartExp t -> SmartExp t
SmartExp (PreSmartExp SmartAcc SmartExp (TAG, ()) -> SmartExp (TAG, ()))
-> PreSmartExp SmartAcc SmartExp (TAG, ()) -> SmartExp (TAG, ())
forall a b. (a -> b) -> a -> b
$ (PreSmartExp SmartAcc SmartExp TAG -> SmartExp TAG
forall t. PreSmartExp SmartAcc SmartExp t -> SmartExp t
SmartExp (PreSmartExp SmartAcc SmartExp TAG -> SmartExp TAG)
-> PreSmartExp SmartAcc SmartExp TAG -> SmartExp TAG
forall a b. (a -> b) -> a -> b
$ PairIdx (TAG, ((), EltR a)) TAG
-> SmartExp (TAG, ((), EltR a))
-> PreSmartExp SmartAcc SmartExp TAG
forall t1 t2 t (exp :: * -> *) (acc :: * -> *).
PairIdx (t1, t2) t -> exp (t1, t2) -> PreSmartExp acc exp t
Prj PairIdx (TAG, ((), EltR a)) TAG
forall a b. PairIdx (a, b) a
PairIdxLeft SmartExp (TAG, ((), EltR a))
SmartExp (EltR (Maybe a))
x) SmartExp TAG
-> SmartExp () -> PreSmartExp SmartAcc SmartExp (TAG, ())
forall (exp :: * -> *) t1 t2 (acc :: * -> *).
exp t1 -> exp t2 -> PreSmartExp acc exp (t1, t2)
`Pair` PreSmartExp SmartAcc SmartExp () -> SmartExp ()
forall t. PreSmartExp SmartAcc SmartExp t -> SmartExp t
SmartExp PreSmartExp SmartAcc SmartExp ()
forall (acc :: * -> *) (exp :: * -> *). PreSmartExp acc exp ()
Nil
  -- TLM: This is a sneaky hack because we know that the tag bits for Just
  -- and True are identical.

-- | The 'fromMaybe' function takes a default value and a 'Maybe' value. If the
-- 'Maybe' is 'Nothing', the default value is returned; otherwise, it returns
-- the value contained in the 'Maybe'.
--
fromMaybe :: Elt a => Exp a -> Exp (Maybe a) -> Exp a
fromMaybe :: Exp a -> Exp (Maybe a) -> Exp a
fromMaybe Exp a
d = (Exp (Maybe a) -> Exp a) -> Exp (Maybe a) -> Exp a
forall f. Matching f => f -> f
match \case
  Exp (Maybe a)
Nothing_ -> Exp a
d
  Just_ Exp a
x  -> Exp a
x

-- | The 'fromJust' function extracts the element out of the 'Just' constructor.
-- If the argument was actually 'Nothing', you will get an undefined value
-- instead.
--
fromJust :: Elt a => Exp (Maybe a) -> Exp a
fromJust :: Exp (Maybe a) -> Exp a
fromJust (Exp SmartExp (EltR (Maybe a))
x) = SmartExp (EltR a) -> Exp a
forall t. SmartExp (EltR t) -> Exp t
Exp (SmartExp (EltR a) -> Exp a) -> SmartExp (EltR a) -> Exp a
forall a b. (a -> b) -> a -> b
$ PreSmartExp SmartAcc SmartExp (EltR a) -> SmartExp (EltR a)
forall t. PreSmartExp SmartAcc SmartExp t -> SmartExp t
SmartExp (PairIdx ((), EltR a) (EltR a)
forall a b. PairIdx (a, b) b
PairIdxRight PairIdx ((), EltR a) (EltR a)
-> SmartExp ((), EltR a) -> PreSmartExp SmartAcc SmartExp (EltR a)
forall t1 t2 t (exp :: * -> *) (acc :: * -> *).
PairIdx (t1, t2) t -> exp (t1, t2) -> PreSmartExp acc exp t
`Prj` PreSmartExp SmartAcc SmartExp ((), EltR a) -> SmartExp ((), EltR a)
forall t. PreSmartExp SmartAcc SmartExp t -> SmartExp t
SmartExp (PairIdx (TAG, ((), EltR a)) ((), EltR a)
forall a b. PairIdx (a, b) b
PairIdxRight PairIdx (TAG, ((), EltR a)) ((), EltR a)
-> SmartExp (TAG, ((), EltR a))
-> PreSmartExp SmartAcc SmartExp ((), EltR a)
forall t1 t2 t (exp :: * -> *) (acc :: * -> *).
PairIdx (t1, t2) t -> exp (t1, t2) -> PreSmartExp acc exp t
`Prj` SmartExp (TAG, ((), EltR a))
SmartExp (EltR (Maybe a))
x))

-- | The 'maybe' function takes a default value, a function, and a 'Maybe'
-- value. If the 'Maybe' value is nothing, the default value is returned;
-- otherwise, it applies the function to the value inside the 'Just' and returns
-- the result
--
maybe :: (Elt a, Elt b) => Exp b -> (Exp a -> Exp b) -> Exp (Maybe a) -> Exp b
maybe :: Exp b -> (Exp a -> Exp b) -> Exp (Maybe a) -> Exp b
maybe Exp b
d Exp a -> Exp b
f = (Exp (Maybe a) -> Exp b) -> Exp (Maybe a) -> Exp b
forall f. Matching f => f -> f
match \case
  Exp (Maybe a)
Nothing_ -> Exp b
d
  Just_ Exp a
x  -> Exp a -> Exp b
f Exp a
x

-- | Extract from an array all of the 'Just' values, together with a segment
-- descriptor indicating how many elements along each dimension were returned.
--
justs :: (Shape sh, Slice sh, Elt a)
      => Acc (Array (sh:.Int) (Maybe a))
      -> Acc (Vector a, Array sh Int)
justs :: Acc (Array (sh :. Int) (Maybe a)) -> Acc (Vector a, Array sh Int)
justs Acc (Array (sh :. Int) (Maybe a))
xs = Acc (Array (sh :. Int) Bool)
-> Acc (Array (sh :. Int) a) -> Acc (Vector a, Array sh Int)
forall sh e.
(Shape sh, Elt e) =>
Acc (Array (sh :. Int) Bool)
-> Acc (Array (sh :. Int) e) -> Acc (Vector e, Array sh Int)
compact ((Exp (Maybe a) -> Exp Bool)
-> Acc (Array (sh :. Int) (Maybe a))
-> Acc (Array (sh :. Int) Bool)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (Maybe a) -> Exp Bool
forall a. Elt a => Exp (Maybe a) -> Exp Bool
isJust Acc (Array (sh :. Int) (Maybe a))
xs) ((Exp (Maybe a) -> Exp a)
-> Acc (Array (sh :. Int) (Maybe a)) -> Acc (Array (sh :. Int) a)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp (Maybe a) -> Exp a
forall a. Elt a => Exp (Maybe a) -> Exp a
fromJust Acc (Array (sh :. Int) (Maybe a))
xs)


instance Functor Maybe where
  fmap :: (Exp a -> Exp b) -> Exp (Maybe a) -> Exp (Maybe b)
fmap Exp a -> Exp b
f = (Exp (Maybe a) -> Exp (Maybe b)) -> Exp (Maybe a) -> Exp (Maybe b)
forall f. Matching f => f -> f
match \case
    Exp (Maybe a)
Nothing_ -> Exp (Maybe b)
forall a. (HasCallStack, Elt a) => Exp (Maybe a)
Nothing_
    Just_ Exp a
x  -> Exp b -> Exp (Maybe b)
forall a. (HasCallStack, Elt a) => Exp a -> Exp (Maybe a)
Just_ (Exp a -> Exp b
f Exp a
x)

instance Eq a => Eq (Maybe a) where
  == :: Exp (Maybe a) -> Exp (Maybe a) -> Exp Bool
(==) = (Exp (Maybe a) -> Exp (Maybe a) -> Exp Bool)
-> Exp (Maybe a) -> Exp (Maybe a) -> Exp Bool
forall f. Matching f => f -> f
match Exp (Maybe a) -> Exp (Maybe a) -> Exp Bool
forall a. Eq a => Exp (Maybe a) -> Exp (Maybe a) -> Exp Bool
go
    where
      go :: Exp (Maybe a) -> Exp (Maybe a) -> Exp Bool
go Exp (Maybe a)
Nothing_  Exp (Maybe a)
Nothing_  = Exp Bool
HasCallStack => Exp Bool
True_
      go (Just_ Exp a
x) (Just_ Exp a
y) = Exp a
x Exp a -> Exp a -> Exp Bool
forall a. Eq a => Exp a -> Exp a -> Exp Bool
== Exp a
y
      go Exp (Maybe a)
_         Exp (Maybe a)
_         = Exp Bool
HasCallStack => Exp Bool
False_

instance Ord a => Ord (Maybe a) where
  compare :: Exp (Maybe a) -> Exp (Maybe a) -> Exp Ordering
compare = (Exp (Maybe a) -> Exp (Maybe a) -> Exp Ordering)
-> Exp (Maybe a) -> Exp (Maybe a) -> Exp Ordering
forall f. Matching f => f -> f
match Exp (Maybe a) -> Exp (Maybe a) -> Exp Ordering
forall a. Ord a => Exp (Maybe a) -> Exp (Maybe a) -> Exp Ordering
go
    where
      go :: Exp (Maybe a) -> Exp (Maybe a) -> Exp Ordering
go (Just_ Exp a
x) (Just_ Exp a
y)  = Exp a -> Exp a -> Exp Ordering
forall a. Ord a => Exp a -> Exp a -> Exp Ordering
compare Exp a
x Exp a
y
      go Exp (Maybe a)
Nothing_  Exp (Maybe a)
Nothing_   = Exp Ordering
HasCallStack => Exp Ordering
EQ_
      go Exp (Maybe a)
Nothing_  Just_{}    = Exp Ordering
HasCallStack => Exp Ordering
LT_
      go Just_{}   Nothing_{} = Exp Ordering
HasCallStack => Exp Ordering
GT_

instance (Monoid (Exp a), Elt a) => Monoid (Exp (Maybe a)) where
  mempty :: Exp (Maybe a)
mempty = Exp (Maybe a)
forall a. (HasCallStack, Elt a) => Exp (Maybe a)
Nothing_

instance (Semigroup (Exp a), Elt a) => Semigroup (Exp (Maybe a)) where
  Exp (Maybe a)
ma <> :: Exp (Maybe a) -> Exp (Maybe a) -> Exp (Maybe a)
<> Exp (Maybe a)
mb = Exp Bool -> Exp (Maybe a) -> Exp (Maybe a) -> Exp (Maybe a)
forall t. Elt t => Exp Bool -> Exp t -> Exp t -> Exp t
cond (Exp (Maybe a) -> Exp Bool
forall a. Elt a => Exp (Maybe a) -> Exp Bool
isNothing Exp (Maybe a)
ma) Exp (Maybe a)
mb
           (Exp (Maybe a) -> Exp (Maybe a)) -> Exp (Maybe a) -> Exp (Maybe a)
forall a b. (a -> b) -> a -> b
$ Exp Bool -> Exp (Maybe a) -> Exp (Maybe a) -> Exp (Maybe a)
forall t. Elt t => Exp Bool -> Exp t -> Exp t -> Exp t
cond (Exp (Maybe a) -> Exp Bool
forall a. Elt a => Exp (Maybe a) -> Exp Bool
isNothing Exp (Maybe a)
mb) Exp (Maybe a)
mb
           (Exp (Maybe a) -> Exp (Maybe a)) -> Exp (Maybe a) -> Exp (Maybe a)
forall a b. (a -> b) -> a -> b
$ Maybe (Exp a) -> Exp (Plain (Maybe (Exp a)))
forall (c :: * -> *) e. Lift c e => e -> c (Plain e)
lift (Exp a -> Maybe (Exp a)
forall a. a -> Maybe a
Just (Exp (Maybe a) -> Exp a
forall a. Elt a => Exp (Maybe a) -> Exp a
fromJust Exp (Maybe a)
ma Exp a -> Exp a -> Exp a
forall a. Semigroup a => a -> a -> a
<> Exp (Maybe a) -> Exp a
forall a. Elt a => Exp (Maybe a) -> Exp a
fromJust Exp (Maybe a)
mb))

instance (Lift Exp a, Elt (Plain a)) => Lift Exp (Maybe a) where
  type Plain (Maybe a) = Maybe (Plain a)
  lift :: Maybe a -> Exp (Plain (Maybe a))
lift Maybe a
Nothing  = Exp (Plain (Maybe a))
forall a. (HasCallStack, Elt a) => Exp (Maybe a)
Nothing_
  lift (Just a
a) = Exp (Plain a) -> Exp (Maybe (Plain a))
forall a. (HasCallStack, Elt a) => Exp a -> Exp (Maybe a)
Just_ (a -> Exp (Plain a)
forall (c :: * -> *) e. Lift c e => e -> c (Plain e)
lift a
a)