Stride.hs

{-# LANGUAGE CPP #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE PatternSynonyms #-}

-- |
-- Module      : Data.Massiv.Core.Index.Stride
-- Copyright   : (c) Alexey Kuleshevich 2018-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
module Data.Massiv.Core.Index.Stride (
  Stride (SafeStride),
  pattern Stride,
  unStride,
  oneStride,
  toLinearIndexStride,
  strideStart,
  strideSize,
) where

import Control.DeepSeq (NFData)
import Data.Massiv.Core.Index.Internal
import System.Random.Stateful (Random, Uniform (..), UniformRange (..))

-- | Stride provides a way to ignore elements of an array if an index is divisible by a
-- corresponding value in a stride. So, for a @Stride (i :. j)@ only elements with indices will be
-- kept around:
--
-- @
-- ( 0 :. 0) ( 0 :. j) ( 0 :. 2j) ( 0 :. 3j) ...
-- ( i :. 0) ( i :. j) ( i :. 2j) ( i :. 3j) ...
-- (2i :. 0) (2i :. j) (2i :. 2j) (2i :. 3j) ...
-- ...
-- @
--
-- Only positive strides make sense, so `Stride` pattern synonym constructor will prevent a user
-- from creating a stride with negative or zero values, thus promoting safety of the library.
--
-- ====__Examples:__
--
-- * Default and minimal stride of @`Stride` (`pureIndex` 1)@ will have no affect and all elements
--   will kept.
--
-- * If stride is @`Stride` 2@, then every 2nd element (i.e. with index 1, 3, 5, ..) will be skipped
--   and only elemnts with indices divisible by 2 will be kept around.
--
-- * In case of two dimensions, if what you want is to keep all rows divisible by 5, but keep every
--   column intact then you'd use @Stride (5 :. 1)@.
--
-- @since 0.2.1
newtype Stride ix = SafeStride ix deriving (Eq, Ord, NFData)

-- | A safe bidirectional pattern synonym for `Stride` construction that will make sure stride
-- elements are always positive.
--
-- @since 0.2.1
pattern Stride :: Index ix => ix -> Stride ix
pattern Stride ix <- SafeStride ix
  where
    Stride ix = SafeStride (liftIndex (max 1) ix)

{-# COMPLETE Stride #-}

instance Index ix => Show (Stride ix) where
  showsPrec n (SafeStride ix) = showsPrecWrapped n (("Stride " ++) . showsPrec 1 ix)

instance (UniformRange ix, Index ix) => Uniform (Stride ix) where
  uniformM g = SafeStride <$> uniformRM (pureIndex 1, pureIndex maxBound) g
  {-# INLINE uniformM #-}

instance UniformRange ix => UniformRange (Stride ix) where
  uniformRM (SafeStride l, SafeStride u) g = SafeStride <$> uniformRM (l, u) g
  {-# INLINE uniformRM #-}

instance (UniformRange ix, Index ix) => Random (Stride ix)

-- | Just a helper function for unwrapping `Stride`.
--
-- @since 0.2.1
unStride :: Stride ix -> ix
unStride (SafeStride ix) = ix
{-# INLINE unStride #-}

-- | Adjust starting index according to the stride
--
-- @since 0.2.1
strideStart :: Index ix => Stride ix -> ix -> ix
strideStart (SafeStride stride) ix =
  liftIndex2
    (+)
    ix
    (liftIndex2 mod (liftIndex2 subtract (liftIndex2 mod ix stride) stride) stride)
{-# INLINE strideStart #-}

-- | Adjust size according to the stride.
--
-- @since 0.2.1
strideSize :: Index ix => Stride ix -> Sz ix -> Sz ix
strideSize (SafeStride stride) (SafeSz sz) =
  SafeSz (liftIndex (+ 1) $ liftIndex2 div (liftIndex (subtract 1) sz) stride)
{-# INLINE strideSize #-}

-- | Compute linear index with stride using the original size and index
--
-- @since 0.2.1
toLinearIndexStride
  :: Index ix
  => Stride ix
  -- ^ Stride
  -> Sz ix
  -- ^ Size
  -> ix
  -- ^ Index
  -> Int
toLinearIndexStride (SafeStride stride) sz ix = toLinearIndex sz (liftIndex2 div ix stride)
{-# INLINE toLinearIndexStride #-}

-- | A default stride of @1@, where all elements are kept
--
-- @since 0.2.1
oneStride :: Index ix => Stride ix
oneStride = SafeStride (pureIndex 1)
{-# INLINE oneStride #-}