{-# Language MagicHash #-}

-- | This solution to holding a sparse set of elements for dynamic programming. The underlying
-- representation requires @O (log n)@ access time for each read or write, where @n@ is the number
-- of elements to be stored. It uses an experimental "bucketing" system to provide better lower and
-- upper bounds than otherwise possible.
--
-- TODO @ADPfusion / FillTyLvl@ handles actually filling the tables. In case all @BigOrder@ tables
-- are dense and of the same dimensional extent, we are fine. However if at least one table is
-- dense, while others are sparse, we will have write to nothing, which should not crash. In case of
-- all-sparse tables for a BigOrder, we have to calculate the union of all indices. This all is
-- currently not happening...
--
-- This version requires working @fromLinearIndex@ but is potentially faster.

module Data.PrimitiveArray.Sparse.IntBinSearch where

import           Control.Monad.Primitive (PrimState,PrimMonad)
import           Control.Monad.ST (ST)
import           Data.Bits.Extras (msb)
import           Debug.Trace (traceShow)
import qualified Control.Monad.State.Strict as SS
import qualified Data.HashMap.Strict as HMS
import qualified Data.Vector.Algorithms.Radix as Sort
import qualified Data.Vector.Algorithms.Search as VAS
import qualified Data.Vector as V
import qualified Data.Vector.Fusion.Stream.Monadic as SM
import qualified Data.Vector.Generic as VG
import qualified Data.Vector.Generic.Mutable as VGM
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Unboxed as VU
import           GHC.Exts ( Int(..), Int#(..), (==#), (-#), (/=#), (*#), (+#), (<=#), remInt#, quotInt#, uncheckedIShiftRA#, (<#) )

import           Data.PrimitiveArray.Class
import           Data.PrimitiveArray.Index.Class
import           Data.PrimitiveArray.Index        -- TODO only while @SparseBucket@ is here



-- | This is a sparse matrix, where only a subset of indices have data associated.

data Sparse w v sh e = Sparse
  { Sparse w v sh e -> LimitType sh
sparseUpperBound  :: !(LimitType sh)
  -- ^ The upper bound for the DP matrix. Not the upper bound of indexes in use, but the theoretical
  -- upper bound.
  , Sparse w v sh e -> v e
sparseData        :: !(v e)
  -- ^ Vector with actually existing data.
  , Sparse w v sh e -> Vector Int
sparseIndices     :: !(VU.Vector Int)
  -- ^ Linearly encoded sparse indices
  , Sparse w v sh e -> Vector Int
manhattanStart    :: !(VU.Vector Int)
  -- ^ Provides left/right boundaries into @sparseIndices@ to speed up index search. Should be one
  -- larger than the largest index to look up, to always provides a good excluded bound.
  }

type Unboxed  w sh e = Sparse w VU.Vector sh e

type Storable w sh e = Sparse w VS.Vector sh e

type Boxed    w sh e = Sparse w  V.Vector sh e



-- | Currently, our mutable variant of sparse matrices will keep indices and manhattan starts
-- immutable as well.

data instance MutArr m (Sparse w v sh e) = MSparse
  { MutArr m (Sparse w v sh e) -> LimitType sh
msparseUpperBound :: !(LimitType sh)
  , MutArr m (Sparse w v sh e) -> Mutable v (PrimState m) e
msparseData       :: !(VG.Mutable v (PrimState m) e)
  , MutArr m (Sparse w v sh e) -> Vector Int
msparseIndices    :: !(VU.Vector Int)
  , MutArr m (Sparse w v sh e) -> Vector Int
mmanhattanStart   :: !(VU.Vector Int)
  }
--  deriving (Generic,Typeable)


type instance FillStruc (Sparse w v sh e) = (w sh)



instance
  ( Index sh, SparseBucket sh, Eq sh, Ord sh
  , VG.Vector w sh , VG.Vector w (Int,sh), VG.Vector w (Int,(Int,sh)), VG.Vector w (Int,Int), VG.Vector w Int
  , VG.Vector v e
#if ADPFUSION_DEBUGOUTPUT
  , Show sh, Show (LimitType sh), Show e
#endif
  ) => PrimArrayOps (Sparse w v) sh e where

  -- ** pure operations

  {-# Inline upperBound #-}
  upperBound :: Sparse w v sh e -> LimitType sh
upperBound Sparse{v e
Vector Int
LimitType sh
manhattanStart :: Vector Int
sparseIndices :: Vector Int
sparseData :: v e
sparseUpperBound :: LimitType sh
manhattanStart :: forall k (w :: k) k (v :: k -> *) sh (e :: k).
Sparse w v sh e -> Vector Int
sparseIndices :: forall k (w :: k) k (v :: k -> *) sh (e :: k).
Sparse w v sh e -> Vector Int
sparseData :: forall k (w :: k) k (v :: k -> *) sh (e :: k).
Sparse w v sh e -> v e
sparseUpperBound :: forall k (w :: k) k (v :: k -> *) sh (e :: k).
Sparse w v sh e -> LimitType sh
..} = LimitType sh
sparseUpperBound
  {-# Inline unsafeIndex #-}
  unsafeIndex :: Sparse w v sh e -> sh -> e
unsafeIndex Sparse{v e
Vector Int
LimitType sh
manhattanStart :: Vector Int
sparseIndices :: Vector Int
sparseData :: v e
sparseUpperBound :: LimitType sh
manhattanStart :: forall k (w :: k) k (v :: k -> *) sh (e :: k).
Sparse w v sh e -> Vector Int
sparseIndices :: forall k (w :: k) k (v :: k -> *) sh (e :: k).
Sparse w v sh e -> Vector Int
sparseData :: forall k (w :: k) k (v :: k -> *) sh (e :: k).
Sparse w v sh e -> v e
sparseUpperBound :: forall k (w :: k) k (v :: k -> *) sh (e :: k).
Sparse w v sh e -> LimitType sh
..} sh
idx = case LimitType sh -> Vector Int -> Vector Int -> sh -> Maybe Int
forall sh.
(SparseBucket sh, Index sh) =>
LimitType sh -> Vector Int -> Vector Int -> sh -> Maybe Int
manhattanIndex LimitType sh
sparseUpperBound Vector Int
manhattanStart Vector Int
sparseIndices sh
idx of
      Maybe Int
Nothing -> [Char] -> e
forall a. HasCallStack => [Char] -> a
error [Char]
"unsafeIndex of non-existing index"
      Just Int
v  -> v e -> Int -> e
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
VG.unsafeIndex v e
sparseData Int
v
  {-# Inline safeIndex #-}
  safeIndex :: Sparse w v sh e -> sh -> Maybe e
safeIndex Sparse{v e
Vector Int
LimitType sh
manhattanStart :: Vector Int
sparseIndices :: Vector Int
sparseData :: v e
sparseUpperBound :: LimitType sh
manhattanStart :: forall k (w :: k) k (v :: k -> *) sh (e :: k).
Sparse w v sh e -> Vector Int
sparseIndices :: forall k (w :: k) k (v :: k -> *) sh (e :: k).
Sparse w v sh e -> Vector Int
sparseData :: forall k (w :: k) k (v :: k -> *) sh (e :: k).
Sparse w v sh e -> v e
sparseUpperBound :: forall k (w :: k) k (v :: k -> *) sh (e :: k).
Sparse w v sh e -> LimitType sh
..} = (Int -> e) -> Maybe Int -> Maybe e
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (v e -> Int -> e
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
VG.unsafeIndex v e
sparseData) (Maybe Int -> Maybe e) -> (sh -> Maybe Int) -> sh -> Maybe e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LimitType sh -> Vector Int -> Vector Int -> sh -> Maybe Int
forall sh.
(SparseBucket sh, Index sh) =>
LimitType sh -> Vector Int -> Vector Int -> sh -> Maybe Int
manhattanIndex LimitType sh
sparseUpperBound Vector Int
manhattanStart Vector Int
sparseIndices

  -- ** monadic operations

  {-# Inline unsafeFreezeM #-}
  unsafeFreezeM :: MutArr m (Sparse w v sh e) -> m (Sparse w v sh e)
unsafeFreezeM MSparse{..} = do
    let sparseUpperBound :: LimitType sh
sparseUpperBound = LimitType sh
msparseUpperBound
        sparseIndices :: Vector Int
sparseIndices = Vector Int
msparseIndices
        manhattanStart :: Vector Int
manhattanStart = Vector Int
mmanhattanStart
    v e
sparseData <- Mutable v (PrimState m) e -> m (v e)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
VG.unsafeFreeze Mutable v (PrimState m) e
msparseData
    Sparse w v sh e -> m (Sparse w v sh e)
forall (m :: * -> *) a. Monad m => a -> m a
return Sparse :: forall k k (w :: k) (v :: k -> *) sh (e :: k).
LimitType sh -> v e -> Vector Int -> Vector Int -> Sparse w v sh e
Sparse{v e
Vector Int
LimitType sh
sparseData :: v e
manhattanStart :: Vector Int
sparseIndices :: Vector Int
sparseUpperBound :: LimitType sh
manhattanStart :: Vector Int
sparseIndices :: Vector Int
sparseData :: v e
sparseUpperBound :: LimitType sh
..}
  {-# Inline unsafeThawM #-}
  unsafeThawM :: Sparse w v sh e -> m (MutArr m (Sparse w v sh e))
unsafeThawM Sparse{v e
Vector Int
LimitType sh
manhattanStart :: Vector Int
sparseIndices :: Vector Int
sparseData :: v e
sparseUpperBound :: LimitType sh
manhattanStart :: forall k (w :: k) k (v :: k -> *) sh (e :: k).
Sparse w v sh e -> Vector Int
sparseIndices :: forall k (w :: k) k (v :: k -> *) sh (e :: k).
Sparse w v sh e -> Vector Int
sparseData :: forall k (w :: k) k (v :: k -> *) sh (e :: k).
Sparse w v sh e -> v e
sparseUpperBound :: forall k (w :: k) k (v :: k -> *) sh (e :: k).
Sparse w v sh e -> LimitType sh
..} = do
    let msparseUpperBound :: LimitType sh
msparseUpperBound = LimitType sh
sparseUpperBound
        msparseIndices :: Vector Int
msparseIndices = Vector Int
sparseIndices
        mmanhattanStart :: Vector Int
mmanhattanStart = Vector Int
manhattanStart
    Mutable v (PrimState m) e
msparseData <- v e -> m (Mutable v (PrimState m) e)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
v a -> m (Mutable v (PrimState m) a)
VG.unsafeThaw v e
sparseData
    MutArr m (Sparse w v sh e) -> m (MutArr m (Sparse w v sh e))
forall (m :: * -> *) a. Monad m => a -> m a
return MSparse :: forall k (m :: * -> *) (w :: k) (v :: * -> *) sh e.
LimitType sh
-> Mutable v (PrimState m) e
-> Vector Int
-> Vector Int
-> MutArr m (Sparse w v sh e)
MSparse{Vector Int
Mutable v (PrimState m) e
LimitType sh
msparseData :: Mutable v (PrimState m) e
mmanhattanStart :: Vector Int
msparseIndices :: Vector Int
msparseUpperBound :: LimitType sh
mmanhattanStart :: Vector Int
msparseIndices :: Vector Int
msparseData :: Mutable v (PrimState m) e
msparseUpperBound :: LimitType sh
..}
  {-# Inline upperBoundM #-}
  upperBoundM :: MutArr m (Sparse w v sh e) -> LimitType sh
upperBoundM MSparse{..} = LimitType sh
msparseUpperBound
  {-# Inline newM #-}
  newM :: LimitType sh -> m (MutArr m (Sparse w v sh e))
newM = [Char] -> LimitType sh -> m (MutArr m (Sparse w v sh e))
forall a. HasCallStack => [Char] -> a
error [Char]
"not implemented, use newSM"
  {-# Inline newWithM #-}
  newWithM :: LimitType sh -> e -> m (MutArr m (Sparse w v sh e))
newWithM = [Char] -> LimitType sh -> e -> m (MutArr m (Sparse w v sh e))
forall a. HasCallStack => [Char] -> a
error [Char]
"not implemented, use newWithSM"
  {-# Inline readM #-}
  readM :: MutArr m (Sparse w v sh e) -> sh -> m e
readM MSparse{..} sh
idx = do
    case LimitType sh -> Vector Int -> Vector Int -> sh -> Maybe Int
forall sh.
(SparseBucket sh, Index sh) =>
LimitType sh -> Vector Int -> Vector Int -> sh -> Maybe Int
manhattanIndex LimitType sh
msparseUpperBound Vector Int
mmanhattanStart Vector Int
msparseIndices sh
idx of
      Maybe Int
Nothing -> [Char] -> m e
forall a. HasCallStack => [Char] -> a
error [Char]
"read of non-existing element"
      Just Int
v  -> Mutable v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead Mutable v (PrimState m) e
msparseData Int
v
  -- | Note that @writeM@ will fail loudly, because we can specialize in @FillTyLvl@ to use
  -- non-failing writes.
  {-# Inline writeM #-}
  writeM :: MutArr m (Sparse w v sh e) -> sh -> e -> m ()
writeM MSparse{..} sh
idx e
elm = do
    case LimitType sh -> Vector Int -> Vector Int -> sh -> Maybe Int
forall sh.
(SparseBucket sh, Index sh) =>
LimitType sh -> Vector Int -> Vector Int -> sh -> Maybe Int
manhattanIndex LimitType sh
msparseUpperBound Vector Int
mmanhattanStart Vector Int
msparseIndices sh
idx of
      Maybe Int
Nothing -> [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"read of non-existing element"
      Just Int
v  -> Mutable v (PrimState m) e -> Int -> e -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite Mutable v (PrimState m) e
msparseData Int
v e
elm
  {-# Inline newSM #-}
  newSM :: LimitType sh
-> FillStruc (Sparse w v sh e) -> m (MutArr m (Sparse w v sh e))
newSM LimitType sh
h FillStruc (Sparse w v sh e)
fs' = do
    let msparseUpperBound :: LimitType sh
msparseUpperBound = LimitType sh
h
        -- sort sparse indices by (manhattan, linearIndex)
        {-# Inline srt #-}
        srt :: Int -> Int -> Ordering
srt Int
x Int
y = let ix :: sh
ix = LimitType sh -> Int -> sh
forall i. Index i => LimitType i -> Int -> i
fromLinearIndex LimitType sh
h Int
x
                      iy :: sh
iy = LimitType sh -> Int -> sh
forall i. Index i => LimitType i -> Int -> i
fromLinearIndex LimitType sh
h Int
y
                  in  (Int, Int) -> (Int, Int) -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (LimitType sh -> sh -> Int
forall sh. SparseBucket sh => LimitType sh -> sh -> Int
manhattan LimitType sh
h sh
ix, Int
x) (LimitType sh -> sh -> Int
forall sh. SparseBucket sh => LimitType sh -> sh -> Int
manhattan LimitType sh
h sh
iy, Int
y)
        {-# Inline radixsrt #-}
        radixsrt :: Int -> Int -> Int
radixsrt Int
i Int
x = let ix :: sh
ix = LimitType sh -> Int -> sh
forall i. Index i => LimitType i -> Int -> i
fromLinearIndex LimitType sh
h Int
x in Int -> (Int, Int) -> Int
forall e. Radix e => Int -> e -> Int
Sort.radix Int
i (LimitType sh -> sh -> Int
forall sh. SparseBucket sh => LimitType sh -> sh -> Int
manhattan LimitType sh
h sh
ix, Int
x)
    Vector Int
msparseIndices <- do
      MVector (PrimState m) Int
marr <- Vector Int -> m (Mutable Vector (PrimState m) Int)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
v a -> m (Mutable v (PrimState m) a)
VG.thaw (w Int -> Vector Int
forall (v :: * -> *) a (w :: * -> *).
(Vector v a, Vector w a) =>
v a -> w a
VU.convert (w Int -> Vector Int) -> w Int -> Vector Int
forall a b. (a -> b) -> a -> b
$ (sh -> Int) -> w sh -> w Int
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
VG.map (LimitType sh -> sh -> Int
forall i. Index i => LimitType i -> i -> Int
linearIndex LimitType sh
h) w sh
FillStruc (Sparse w v sh e)
fs')
      Int
-> Int -> (Int -> Int -> Int) -> MVector (PrimState m) Int -> m ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Int -> Int -> (Int -> e -> Int) -> v (PrimState m) e -> m ()
Sort.sortBy ((Int, Int) -> Int
forall e. Radix e => e -> Int
Sort.passes ((Int, Int)
forall a. HasCallStack => a
undefined :: (Int,Int))) (Int -> Int
forall e. Radix e => e -> Int
Sort.size (Int
forall a. HasCallStack => a
undefined :: Int)) Int -> Int -> Int
radixsrt MVector (PrimState m) Int
marr
      Mutable Vector (PrimState m) Int -> m (Vector Int)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
Mutable v (PrimState m) a -> m (v a)
VG.unsafeFreeze MVector (PrimState m) Int
Mutable Vector (PrimState m) Int
marr
    let -- For any manhattan distance not found in the distances, we set to the length of the the
        -- @msparseIndices@ vector. Perform reverse-scan to update all manhattan start distances.
        go :: VU.MVector s Int -> ST s ()
        {-# Inline go #-}
        go :: MVector s Int -> ST s ()
go MVector s Int
mv = do
          Vector (Int, Int) -> ((Int, Int) -> ST s ()) -> ST s ()
forall (m :: * -> *) (v :: * -> *) a b.
(Monad m, Vector v a) =>
v a -> (a -> m b) -> m ()
VG.forM_ (Vector (Int, Int) -> Vector (Int, Int)
forall (v :: * -> *) a. Vector v a => v a -> v a
VG.reverse (Vector (Int, Int) -> Vector (Int, Int))
-> Vector (Int, Int) -> Vector (Int, Int)
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector (Int, Int)
forall (v :: * -> *) a.
(Vector v a, Vector v (Int, a)) =>
v a -> v (Int, a)
VG.indexed Vector Int
msparseIndices) (((Int, Int) -> ST s ()) -> ST s ())
-> ((Int, Int) -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \(Int
i,Int
k) -> let lix :: sh
lix = LimitType sh -> Int -> sh
forall i. Index i => LimitType i -> Int -> i
fromLinearIndex LimitType sh
h Int
k; mh :: Int
mh = LimitType sh -> sh -> Int
forall sh. SparseBucket sh => LimitType sh -> sh -> Int
manhattan LimitType sh
h sh
lix in MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector s Int
MVector (PrimState (ST s)) Int
mv Int
mh Int
i
    let mmanhattanStart :: Vector Int
mmanhattanStart = (forall s. MVector s Int -> ST s ()) -> Vector Int -> Vector Int
forall a.
Unbox a =>
(forall s. MVector s a -> ST s ()) -> Vector a -> Vector a
VU.modify forall s. MVector s Int -> ST s ()
go (Vector Int -> Vector Int) -> Vector Int -> Vector Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Vector Int
forall (v :: * -> *) a. Vector v a => Int -> a -> v a
VG.replicate (LimitType sh -> Int
forall sh. SparseBucket sh => LimitType sh -> Int
manhattanMax LimitType sh
h Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Vector Int -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector Int
msparseIndices)
    Mutable v (PrimState m) e
msparseData <- Int -> m (Mutable v (PrimState m) e)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
VGM.new (Int -> m (Mutable v (PrimState m) e))
-> Int -> m (Mutable v (PrimState m) e)
forall a b. (a -> b) -> a -> b
$ Vector Int -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector Int
msparseIndices
    MutArr m (Sparse w v sh e) -> m (MutArr m (Sparse w v sh e))
forall (m :: * -> *) a. Monad m => a -> m a
return (MutArr m (Sparse w v sh e) -> m (MutArr m (Sparse w v sh e)))
-> MutArr m (Sparse w v sh e) -> m (MutArr m (Sparse w v sh e))
forall a b. (a -> b) -> a -> b
$ MSparse :: forall k (m :: * -> *) (w :: k) (v :: * -> *) sh e.
LimitType sh
-> Mutable v (PrimState m) e
-> Vector Int
-> Vector Int
-> MutArr m (Sparse w v sh e)
MSparse {Vector Int
Mutable v (PrimState m) e
LimitType sh
msparseData :: Mutable v (PrimState m) e
mmanhattanStart :: Vector Int
msparseIndices :: Vector Int
msparseUpperBound :: LimitType sh
mmanhattanStart :: Vector Int
msparseIndices :: Vector Int
msparseData :: Mutable v (PrimState m) e
msparseUpperBound :: LimitType sh
..}
  {-# Inline newWithSM #-}
  newWithSM :: LimitType sh
-> FillStruc (Sparse w v sh e)
-> e
-> m (MutArr m (Sparse w v sh e))
newWithSM LimitType sh
h FillStruc (Sparse w v sh e)
fs' e
e = do
    MutArr m (Sparse w v sh e)
mv <- LimitType sh
-> FillStruc (Sparse w v sh e) -> m (MutArr m (Sparse w v sh e))
forall (arr :: * -> * -> *) sh elm (m :: * -> *).
(PrimArrayOps arr sh elm, Monad m, PrimMonad m) =>
LimitType sh -> FillStruc (arr sh elm) -> m (MutArr m (arr sh elm))
newSM LimitType sh
h FillStruc (Sparse w v sh e)
fs'
    Mutable v (PrimState m) e -> e -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> a -> m ()
VGM.set (MutArr m (Sparse w v sh e) -> Mutable v (PrimState m) e
forall (m :: * -> *) k (w :: k) (v :: * -> *) sh e.
MutArr m (Sparse w v sh e) -> Mutable v (PrimState m) e
msparseData MutArr m (Sparse w v sh e)
mv) e
e
    MutArr m (Sparse w v sh e) -> m (MutArr m (Sparse w v sh e))
forall (m :: * -> *) a. Monad m => a -> m a
return MutArr m (Sparse w v sh e)
mv
  {-# Inline safeWriteM #-}
  safeWriteM :: MutArr m (Sparse w v sh e) -> sh -> e -> m ()
safeWriteM MSparse{..} sh
sh e
e = case LimitType sh -> Vector Int -> Vector Int -> sh -> Maybe Int
forall sh.
(SparseBucket sh, Index sh) =>
LimitType sh -> Vector Int -> Vector Int -> sh -> Maybe Int
manhattanIndex LimitType sh
msparseUpperBound Vector Int
mmanhattanStart Vector Int
msparseIndices sh
sh of
      Maybe Int
Nothing -> () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
      Just Int
v  -> Mutable v (PrimState m) e -> Int -> e -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite Mutable v (PrimState m) e
msparseData Int
v e
e
  {-# Inline safeReadM #-}
  safeReadM :: MutArr m (Sparse w v sh e) -> sh -> m (Maybe e)
safeReadM MSparse{..} sh
sh = case LimitType sh -> Vector Int -> Vector Int -> sh -> Maybe Int
forall sh.
(SparseBucket sh, Index sh) =>
LimitType sh -> Vector Int -> Vector Int -> sh -> Maybe Int
manhattanIndex LimitType sh
msparseUpperBound Vector Int
mmanhattanStart Vector Int
msparseIndices sh
sh of
      Maybe Int
Nothing -> Maybe e -> m (Maybe e)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe e
forall a. Maybe a
Nothing
      Just Int
v  -> e -> Maybe e
forall a. a -> Maybe a
Just (e -> Maybe e) -> m e -> m (Maybe e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Mutable v (PrimState m) e -> Int -> m e
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
VGM.unsafeRead Mutable v (PrimState m) e
msparseData Int
v
  -- ** implement me
  transformShape :: (LimitType sh -> LimitType sh')
-> Sparse w v sh e -> Sparse w v sh' e
transformShape = [Char]
-> (LimitType sh -> LimitType sh')
-> Sparse w v sh e
-> Sparse w v sh' e
forall a. HasCallStack => [Char] -> a
error [Char]
"implement me"
  fromListM :: LimitType sh -> [e] -> m (MutArr m (Sparse w v sh e))
fromListM = [Char] -> LimitType sh -> [e] -> m (MutArr m (Sparse w v sh e))
forall a. HasCallStack => [Char] -> a
error [Char]
"implement me"



instance (Index sh, VG.Vector v e, VG.Vector v e') ⇒ PrimArrayMap (Sparse w v) sh e e' where
  {-# Inline mapArray #-}
  mapArray :: (e -> e') -> Sparse w v sh e -> Sparse w v sh e'
mapArray e -> e'
f Sparse w v sh e
sparse = Sparse w v sh e
sparse{sparseData :: v e'
sparseData = (e -> e') -> v e -> v e'
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
VG.map e -> e'
f (Sparse w v sh e -> v e
forall k (w :: k) k (v :: k -> *) sh (e :: k).
Sparse w v sh e -> v e
sparseData Sparse w v sh e
sparse)}





-- * Helper functions.

-- | Find the index with manhattan helper
--
-- TODO consider using binary search instead of a linear scan here!
-- e.g.: @k = VAS.binarySearchByBounds (==)@
--
-- NOTE running times with 100x100 DP problem "NeedlemanWunsch"
-- full findIndex of sixs:              0,050,000 cells/sec
-- using manhattan buckets, findIndex:  5,000,000 cells/sec
-- using binarySearch on slices:       11,000,000 cells/sec
--
-- On a 1000x1000 DP NeedlemanWunsch problem, binary search on slices is at 6,500,000 cells/sec.

manhattanIndex
  :: (SparseBucket sh, Index sh)
  => LimitType sh -> Vector Int -> VU.Vector Int -> sh -> Maybe Int
{-# Inline manhattanIndex #-}
manhattanIndex :: LimitType sh -> Vector Int -> Vector Int -> sh -> Maybe Int
manhattanIndex LimitType sh
ub Vector Int
mstart Vector Int
sixs sh
idx = (Int -> Int) -> Maybe Int -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
l) (Maybe Int -> Maybe Int)
-> (Vector Int -> Maybe Int) -> Vector Int -> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Vector Int -> Maybe Int
binarySearch (LimitType sh -> sh -> Int
forall i. Index i => LimitType i -> i -> Int
linearIndex LimitType sh
ub sh
idx) (Vector Int -> Maybe Int) -> Vector Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Vector Int -> Vector Int
forall (v :: * -> *) a. Vector v a => Int -> Int -> v a -> v a
VG.unsafeSlice Int
l (Int
hInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
l) Vector Int
sixs
  where
    b :: Int
b = LimitType sh -> sh -> Int
forall sh. SparseBucket sh => LimitType sh -> sh -> Int
manhattan LimitType sh
ub sh
idx
    -- lower and upper bucket bounds
    l :: Int
l = Vector Int
mstart Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
`VU.unsafeIndex` Int
b
    h :: Int
h = Vector Int
mstart Vector Int -> Int -> Int
forall a. Unbox a => Vector a -> Int -> a
`VU.unsafeIndex` (Int
bInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

binarySearch :: Int -> VU.Vector Int -> Maybe Int
{-# Inline binarySearch #-}
{-
binarySearch k xs =
  let r1 = binarySearch1 k xs
      r2 = binarySearch2 k xs
  in if r1==r2 then r1 else error $ show (k,xs,r1,r2)
-}
{-
-- 1000x1000 at @1000 yields 3,050,000 cells / second
binarySearch (I# e) v = go 0 pp
  where
    -- largest index to check
    (I# r) = VG.length v -1
    -- largest power of two <= (r+1)
    pp = (2 ^ (max 0 $ msb $ VG.length v -1))
    -- wrap the actual non-branching worker function
    go :: Int -> Int -> Maybe Int
    {-# Inline go #-}
    go (I# l) (I# p) = let i = I# (go' l p) in if (VG.length v<1 || i<0) then Nothing else Just i
    -- @go'@ should be non-branching, and use a minimal number of array reads.
    go' :: Int# -> Int# -> Int#
    {-# Inline go' #-}
    go' l p
      -- we are done and will return the proposed position of the last element found or -1
      | 1# <- p ==# 0# = (e ==# x) *# l -# (e /=# x)
      | otherwise      = let i = go' (l +# (p *# chk *# leq)) (quotInt# p 2#) -- (uncheckedIShiftRA# p 1#)
                             -- (I# ii) = traceShow (I# l, I# p, I# i2r, I# x, I# leq) (I# i)
                         in  i -- i
      where i2r    = l +# (p *# chk) -- index to read
            (I# x) = VU.unsafeIndex v (I# i2r)
            leq    = x <=# e
            newl   = l +# p
            chk    = newl <=# r
-}
--
-- 1000x1000 at @1000 yields 6,030,000 cells / second
binarySearch :: Int -> Vector Int -> Maybe Int
binarySearch Int
e Vector Int
v = Int -> Int -> Maybe Int
go Int
0 (Vector Int -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length Vector Int
v Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
  where
    go :: Int -> Int -> Maybe Int
    go :: Int -> Int -> Maybe Int
go !Int
l !Int
r =
      let !m :: Int
m = (Int
rInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
l) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
          !x :: Int
x = Vector Int -> Int -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int -> a
VG.unsafeIndex Vector Int
v Int
m
      in  if Int
rInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
l then Maybe Int
forall a. Maybe a
Nothing else case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
e Int
x of
                                    Ordering
LT -> Int -> Int -> Maybe Int
go Int
l (Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
                                    Ordering
EQ -> Int -> Maybe Int
forall a. a -> Maybe a
Just Int
m
                                    Ordering
GT -> Int -> Int -> Maybe Int
go (Int
mInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
r
--


-- | Given two index vectors of the same shape, will return the correctly ordered vector of
-- the union of indices.
--
-- TODO This requires that @Ord (Shape O)@ uses the @Down@ instance of Ord! We need to fix this in
-- the @Index@ modules.
--
-- TODO Rewrite to allow fusion without intermediate vectors using uncons. This will make it
-- possible to chain applications. @stream@ should be fine for this.

mergeIndexVectors :: (Eq sh, Ord sh, VG.Vector w sh) => w sh -> w sh -> w sh
{-# Inlinable mergeIndexVectors #-}
mergeIndexVectors :: w sh -> w sh -> w sh
mergeIndexVectors w sh
xs w sh
ys = (forall s. ST s (Mutable w s sh)) -> w sh
forall (v :: * -> *) a.
Vector v a =>
(forall s. ST s (Mutable v s a)) -> v a
VG.create ((forall s. ST s (Mutable w s sh)) -> w sh)
-> (forall s. ST s (Mutable w s sh)) -> w sh
forall a b. (a -> b) -> a -> b
$ do
  let lxs :: Int
lxs = w sh -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length w sh
xs
      lys :: Int
lys = w sh -> Int
forall (v :: * -> *) a. Vector v a => v a -> Int
VG.length w sh
ys
  Mutable w s sh
mv <- Int -> ST s (Mutable w (PrimState (ST s)) sh)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> m (v (PrimState m) a)
VGM.new (Int -> ST s (Mutable w (PrimState (ST s)) sh))
-> Int -> ST s (Mutable w (PrimState (ST s)) sh)
forall a b. (a -> b) -> a -> b
$ Int
lxs Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
lys
  let go :: Int -> Int -> Int -> ST s Int
go !Int
n !Int
i !Int
j
        | Int
iInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>=Int
lxs Bool -> Bool -> Bool
&& Int
jInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>=Int
lys = Int -> ST s Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
n
        | Int
jInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>=Int
lys = w sh -> Int -> ST s sh
forall (v :: * -> *) a (m :: * -> *).
(Vector v a, Monad m) =>
v a -> Int -> m a
VG.unsafeIndexM w sh
xs Int
i ST s sh -> (sh -> ST s ()) -> ST s ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Mutable w (PrimState (ST s)) sh -> Int -> sh -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite Mutable w s sh
Mutable w (PrimState (ST s)) sh
mv Int
n ST s () -> ST s Int -> ST s Int
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> Int -> ST s Int
go (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
j
        | Int
iInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>=Int
lxs = w sh -> Int -> ST s sh
forall (v :: * -> *) a (m :: * -> *).
(Vector v a, Monad m) =>
v a -> Int -> m a
VG.unsafeIndexM w sh
ys Int
j ST s sh -> (sh -> ST s ()) -> ST s ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Mutable w (PrimState (ST s)) sh -> Int -> sh -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite Mutable w s sh
Mutable w (PrimState (ST s)) sh
mv Int
n ST s () -> ST s Int -> ST s Int
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> Int -> ST s Int
go (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
i (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
        | Bool
otherwise = do
            sh
x <- w sh -> Int -> ST s sh
forall (v :: * -> *) a (m :: * -> *).
(Vector v a, Monad m) =>
v a -> Int -> m a
VG.unsafeIndexM w sh
xs Int
i
            sh
y <- w sh -> Int -> ST s sh
forall (v :: * -> *) a (m :: * -> *).
(Vector v a, Monad m) =>
v a -> Int -> m a
VG.unsafeIndexM w sh
ys Int
j
            if | sh
xsh -> sh -> Bool
forall a. Eq a => a -> a -> Bool
==sh
y -> Mutable w (PrimState (ST s)) sh -> Int -> sh -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite Mutable w s sh
Mutable w (PrimState (ST s)) sh
mv Int
n sh
x ST s () -> ST s Int -> ST s Int
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> Int -> ST s Int
go (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
               | sh
xsh -> sh -> Bool
forall a. Ord a => a -> a -> Bool
< sh
y -> Mutable w (PrimState (ST s)) sh -> Int -> sh -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite Mutable w s sh
Mutable w (PrimState (ST s)) sh
mv Int
n sh
x ST s () -> ST s Int -> ST s Int
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> Int -> ST s Int
go (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
j
               | sh
xsh -> sh -> Bool
forall a. Ord a => a -> a -> Bool
> sh
y -> Mutable w (PrimState (ST s)) sh -> Int -> sh -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.unsafeWrite Mutable w s sh
Mutable w (PrimState (ST s)) sh
mv Int
n sh
y ST s () -> ST s Int -> ST s Int
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> Int -> Int -> ST s Int
go (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
i     (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
  Int
n <- Int -> Int -> Int -> ST s Int
go Int
0 Int
0 Int
0
  Mutable w s sh -> ST s (Mutable w s sh)
forall (m :: * -> *) a. Monad m => a -> m a
return (Mutable w s sh -> ST s (Mutable w s sh))
-> Mutable w s sh -> ST s (Mutable w s sh)
forall a b. (a -> b) -> a -> b
$ Int -> Mutable w s sh -> Mutable w s sh
forall (v :: * -> * -> *) a s. MVector v a => Int -> v s a -> v s a
VGM.unsafeTake Int
n Mutable w s sh
mv