Pull.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}

-- |
-- Module      : Data.Massiv.Array.Delayed.Pull
-- Copyright   : (c) Alexey Kuleshevich 2018-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
module Data.Massiv.Array.Delayed.Pull (
  D (..),
  Array (..),
  delay,
  eqArrays,
  compareArrays,
  imap,
  liftArray2',
  liftArray2M,
  unsafeExtract,
  unsafeSlice,
  unsafeInnerSlice,
  zipWithInternal,
) where

import Control.Applicative
import qualified Data.Foldable as F
import Data.Massiv.Array.Ops.Fold.Internal as A
import Data.Massiv.Core.Common as A
import Data.Massiv.Core.List (L, showArrayList, showsArrayPrec)
import Data.Massiv.Core.Operations
import qualified Data.Massiv.Vector.Stream as S
import GHC.Base (build)
import Prelude hiding (zipWith)

#include "massiv.h"

-- | Delayed representation.
data D
  = D
  deriving (Show)

data instance Array D ix e = DArray
  { dComp :: !Comp
  , dSize :: !(Sz ix)
  , dPrefIndex :: !(PrefIndex ix e)
  }

instance (Ragged L ix e, Show e) => Show (Array D ix e) where
  showsPrec = showsArrayPrec id
  showList = showArrayList

instance Index ix => Shape D ix where
  maxLinearSize = Just . SafeSz . elemsCount
  {-# INLINE maxLinearSize #-}

instance Size D where
  size = dSize
  {-# INLINE size #-}
  unsafeResize !sz !arr =
    makeArrayLinear (dComp arr) sz (unsafeIndex arr . fromLinearIndex (size arr))
  {-# INLINE unsafeResize #-}

instance Strategy D where
  setComp c arr = arr{dComp = c}
  {-# INLINE setComp #-}
  getComp = dComp
  {-# INLINE getComp #-}
  repr = D

instance Source D e where
  unsafeIndex arr =
    case dPrefIndex arr of
      PrefIndex f -> f
      PrefIndexLinear f -> f . toLinearIndex (size arr)
  {-# INLINE unsafeIndex #-}
  unsafeLinearIndex arr =
    case dPrefIndex arr of
      PrefIndex f -> f . fromLinearIndex (size arr)
      PrefIndexLinear f -> f
  {-# INLINE unsafeLinearIndex #-}
  unsafePrefIndex = dPrefIndex
  {-# INLINE unsafePrefIndex #-}
  unsafeOuterSlice !arr !szL !i =
    makeArray (dComp arr) szL (unsafeIndex arr . consDim i)
  {-# INLINE unsafeOuterSlice #-}
  unsafeLinearSlice !o !sz arr =
    makeArrayLinear (dComp arr) sz $ \ !i -> unsafeLinearIndex arr (i + o)
  {-# INLINE unsafeLinearSlice #-}

-- | /O(1)/ - Extract a portion of an array. Staring index and new size are
-- not validated.
unsafeExtract :: (Source r e, Index ix) => ix -> Sz ix -> Array r ix e -> Array D ix e
unsafeExtract !sIx !newSz !arr =
  makeArray (getComp arr) newSz (unsafeIndex arr . liftIndex2 (+) sIx)
{-# INLINE unsafeExtract #-}

-- | /O(1)/ - Take a slice out of an array from within
unsafeSlice
  :: (Source r e, Index ix, Index (Lower ix), MonadThrow m)
  => Array r ix e
  -> ix
  -> Sz ix
  -> Dim
  -> m (Array D (Lower ix) e)
unsafeSlice arr start cut@(SafeSz cutSz) dim = do
  newSz <- dropDimM cutSz dim
  return $ unsafeResize (SafeSz newSz) (unsafeExtract start cut arr)
{-# INLINE unsafeSlice #-}

-- | /O(1)/ - Take a slice out of an array from the inside
unsafeInnerSlice
  :: (Source r e, Index ix) => Array r ix e -> Sz (Lower ix) -> Int -> Array D (Lower ix) e
unsafeInnerSlice !arr szL !i =
  DArray (getComp arr) szL $ PrefIndex (unsafeIndex arr . (`snocDim` i))
{-# INLINE unsafeInnerSlice #-}

instance (Eq e, Index ix) => Eq (Array D ix e) where
  (==) = eqArrays (==)
  {-# INLINE (==) #-}

instance (Ord e, Index ix) => Ord (Array D ix e) where
  compare = compareArrays compare
  {-# INLINE compare #-}

instance Functor (Array D ix) where
  fmap f (DArray c sz g) = DArray c sz (fmap f g)
  {-# INLINE fmap #-}
  (<$) e (DArray c sz g) = DArray c sz (e <$ g)
  {-# INLINE (<$) #-}

instance Index ix => Applicative (Array D ix) where
  pure = singleton
  {-# INLINE pure #-}
  (<*>) = liftArray2' id
  {-# INLINE (<*>) #-}
#if MIN_VERSION_base(4,10,0)
  liftA2 = liftArray2'
  {-# INLINE liftA2 #-}
#endif

-- | Row-major sequential folding over a Delayed array.
instance Index ix => Foldable (Array D ix) where
  fold = A.fold
  {-# INLINE fold #-}
  foldMap = A.foldMono
  {-# INLINE foldMap #-}
  foldl = lazyFoldlS
  {-# INLINE foldl #-}
  foldl' = foldlS
  {-# INLINE foldl' #-}
  foldr = foldrFB
  {-# INLINE foldr #-}
  foldr' = foldrS
  {-# INLINE foldr' #-}
  null (DArray _ sz _) = totalElem sz == 0
  {-# INLINE null #-}
  length = totalElem . size
  {-# INLINE length #-}
  elem e = A.any (e ==)
  {-# INLINE elem #-}
  toList arr = build (\c n -> foldrFB c n arr)
  {-# INLINE toList #-}

instance Index ix => Load D ix e where
  makeArray comp sz = DArray comp sz . PrefIndex
  {-# INLINE makeArray #-}
  makeArrayLinear comp sz = DArray comp sz . PrefIndexLinear
  {-# INLINE makeArrayLinear #-}
  iterArrayLinearST_ !scheduler DArray{..} uWrite =
    case dPrefIndex of
      PrefIndex f ->
        iterTargetFullST_ defRowMajor scheduler 0 dSize $ \ !i -> uWrite i . f
      PrefIndexLinear f ->
        iterTargetFullST_ defRowMajorLinear scheduler 0 dSize $ \ !i _ -> uWrite i (f i)
  {-# INLINE iterArrayLinearST_ #-}

instance Index ix => StrideLoad D ix e where
  iterArrayLinearWithStrideST_ !scheduler !stride sz DArray{..} uWrite =
    case dPrefIndex of
      PrefIndex f ->
        iterTargetFullWithStrideST_ defRowMajor scheduler 0 sz stride $ \i ->
          uWrite i . f
      PrefIndexLinear f -> do
        iterTargetFullWithStrideST_ defRowMajor scheduler 0 sz stride $ \i ->
          uWrite i . f . toLinearIndex dSize
  {-# INLINE iterArrayLinearWithStrideST_ #-}

instance Index ix => Stream D ix e where
  toStream = S.steps
  {-# INLINE toStream #-}
  toStreamIx = S.steps . imap (,)
  {-# INLINE toStreamIx #-}

-- | Map an index aware function over an array
--
-- @since 0.1.0
imap
  :: forall r ix e a
   . (Index ix, Source r e)
  => (ix -> e -> a)
  -> Array r ix e
  -> Array D ix a
imap f !arr =
  case unsafePrefIndex arr of
    PrefIndex gix -> DArray (getComp arr) sz $ PrefIndex (\ !ix -> f ix (gix ix))
    PrefIndexLinear gi ->
      DArray (getComp arr) sz $ PrefIndex (\ !ix -> f ix (gi (toLinearIndex sz ix)))
  where
    !sz = size arr
{-# INLINE imap #-}

instance Num e => FoldNumeric D e where
  unsafeDotProduct = defaultUnsafeDotProduct
  {-# INLINE unsafeDotProduct #-}
  powerSumArray = defaultPowerSumArray
  {-# INLINE powerSumArray #-}
  foldArray = defaultFoldArray
  {-# INLINE foldArray #-}

instance Num e => Numeric D e where
  unsafeLiftArray f arr = arr{dPrefIndex = f <$> dPrefIndex arr}
  {-# INLINE unsafeLiftArray #-}
  unsafeLiftArray2 f a1 a2 = zipWithInternal (size a1) f a1 a2
  {-# INLINE unsafeLiftArray2 #-}

instance Floating e => NumericFloat D e

-- | /O(1)/ Conversion from a source array to `D` representation.
delay :: (Index ix, Source r e) => Array r ix e -> Array D ix e
delay arr =
  case unsafePrefIndex arr of
    PrefIndex gix -> makeArray (getComp arr) (size arr) gix
    PrefIndexLinear gi -> makeArrayLinear (getComp arr) (size arr) gi
{-# INLINE [1] delay #-}

{-# RULES
"delay" [~1] forall (arr :: Array D ix e). delay arr = arr
  #-}

-- | Compute array equality by applying a comparing function to each
-- element. Empty arrays are always equal, regardless of their size.
--
-- @since 0.5.7
eqArrays
  :: (Index ix, Source r1 e1, Source r2 e2)
  => (e1 -> e2 -> Bool)
  -> Array r1 ix e1
  -> Array r2 ix e2
  -> Bool
eqArrays f arr1 arr2 =
  let sz1 = size arr1
      sz2 = size arr2
   in ( sz1 == sz2
          && not
            ( A.any
                not
                ( makeArray @D (getComp arr1 <> getComp arr2) (size arr1) $ \ix ->
                    f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)
                )
            )
      )
        || (isZeroSz sz1 && isZeroSz sz2)
{-# INLINE eqArrays #-}

-- | Compute array ordering by applying a comparing function to each element.
-- The exact ordering is unspecified so this is only intended for use in maps and the like where
-- you need an ordering but do not care about which one is used.
--
-- @since 0.5.7
compareArrays
  :: (Index ix, Source r1 e1, Source r2 e2)
  => (e1 -> e2 -> Ordering)
  -> Array r1 ix e1
  -> Array r2 ix e2
  -> Ordering
compareArrays f arr1 arr2 =
  compare (size arr1) (size arr2)
    <> A.fold
      ( makeArray @D (getComp arr1 <> getComp arr2) (size arr1) $ \ix ->
          f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix)
      )
{-# INLINE compareArrays #-}

-- | Same as `liftArray2M`, but throws an imprecise exception on mismatched
-- sizes.
--
-- @since 1.0.0
liftArray2'
  :: (HasCallStack, Index ix, Source r1 a, Source r2 b)
  => (a -> b -> e)
  -> Array r1 ix a
  -> Array r2 ix b
  -> Array D ix e
liftArray2' f arr1 arr2 = throwEither $ liftArray2M f arr1 arr2
{-# INLINE liftArray2' #-}

-- | Similar to `Data.Massiv.Array.zipWith`, except dimensions of both arrays
-- have to be the same, otherwise it throws `SizeMismatchException`.
--
-- @since 1.0.0
liftArray2M
  :: (Index ix, Source r1 a, Source r2 b, MonadThrow m)
  => (a -> b -> e)
  -> Array r1 ix a
  -> Array r2 ix b
  -> m (Array D ix e)
liftArray2M f !arr1 !arr2
  | sz1 == sz2 = pure $ zipWithInternal sz1 f arr1 arr2
  | isZeroSz sz1 && isZeroSz sz2 = pure A.empty
  | otherwise = throwM $ SizeMismatchException (size arr1) (size arr2)
  where
    sz1 = size arr1
    sz2 = size arr2
{-# INLINE liftArray2M #-}

zipWithInternal
  :: (Index ix, Source r1 e1, Source r2 e2)
  => Sz ix
  -> (e1 -> e2 -> e3)
  -> Array r1 ix e1
  -> Array r2 ix e2
  -> Array D ix e3
zipWithInternal sz f arr1 arr2 =
  case unsafePrefIndex arr1 of
    PrefIndexLinear gi1
      | PrefIndexLinear gi2 <- unsafePrefIndex arr2 ->
          makeArrayLinear comp sz (\ !i -> f (gi1 i) (gi2 i))
    _ -> makeArray comp sz (\ !ix -> f (unsafeIndex arr1 ix) (unsafeIndex arr2 ix))
  where
    comp = getComp arr1 <> getComp arr2
{-# INLINE zipWithInternal #-}