List.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
-- |
-- Module      : Data.Massiv.Core.List
-- Copyright   : (c) Alexey Kuleshevich 2018-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Data.Massiv.Core.List
  ( L(..)
  , Array(..)
  , List(..)
  , toListArray
  , showsArrayPrec
  , showArrayList
  , ListItem
  ) where

import Control.Monad (unless, when)
import Control.Scheduler
import Data.Coerce
import Data.Functor.Identity
import Data.Kind
import qualified Data.List as L
import Data.Massiv.Core.Common
import qualified Data.Massiv.Vector.Stream as S
import Data.Monoid
import Data.Typeable
import GHC.Exts
import GHC.TypeLits
import System.IO.Unsafe (unsafePerformIO)


type family ListItem ix e :: Type where
  ListItem Ix1 e = e
  ListItem ix  e = [ListItem (Lower ix) e]

type family Elt ix e :: Type where
  Elt Ix1 e = e
  Elt ix  e = List (Lower ix) e

newtype List ix e = List { unList :: [Elt ix e] }


instance Coercible (Elt ix e) (ListItem ix e) => IsList (List ix e) where
  type Item (List ix e) = ListItem ix e
  fromList = coerce
  {-# INLINE fromList #-}
  toList = coerce
  {-# INLINE toList #-}


data L = L

data instance Array L ix e = LArray { lComp :: Comp
                                    , lData :: !(List ix e)
                                    }


instance Coercible (Elt ix e) (ListItem ix e) => IsList (Array L ix e) where
  type Item (Array L ix e) = ListItem ix e
  fromList = LArray Seq . coerce
  {-# INLINE fromList #-}
  toList = coerce . lData
  {-# INLINE toList #-}

lengthHintList :: [a] -> LengthHint
lengthHintList =
  \case
    [] -> LengthExact zeroSz
    _  -> LengthUnknown
{-# INLINE lengthHintList #-}

instance Shape L Ix1 where
  linearSize = outerLength
  {-# INLINE linearSize #-}
  linearSizeHint = lengthHintList . unList . lData
  {-# INLINE linearSizeHint #-}
  isNull = null . unList . lData
  {-# INLINE isNull #-}
  outerSize = linearSize
  {-# INLINE outerSize #-}

instance Shape L Ix2 where
  linearSize = SafeSz . getSum . foldMap (Sum . length . unList) . unList . lData
  {-# INLINE linearSize #-}
  linearSizeHint = lengthHintList . unList . lData
  {-# INLINE linearSizeHint #-}
  isNull = getAll . foldMap (All . null . unList) . unList . lData
  {-# INLINE isNull #-}
  outerSize arr =
    case unList (lData arr) of
      []     -> zeroSz
      (x:xs) -> SafeSz ((1 + length xs) :. length (unList x))
  {-# INLINE outerSize #-}

instance (Shape L (Ix (n - 1)), Index (IxN n)) => Shape L (IxN n) where
  linearSize = SafeSz . getSum . foldMap (Sum . unSz . linearSize . LArray Seq) . unList . lData
  {-# INLINE linearSize #-}
  linearSizeHint = lengthHintList . unList . lData
  {-# INLINE linearSizeHint #-}
  isNull = getAll . foldMap (All . isNull . LArray Seq) . unList . lData
  {-# INLINE isNull #-}
  outerSize arr =
    case unList (lData arr) of
      []     -> zeroSz
      (x:xs) -> SafeSz ((1 + length xs) :> unSz (outerSize (LArray Seq x)))
  {-# INLINE outerSize #-}


outerLength :: Array L ix e -> Sz Int
outerLength = SafeSz . length . unList . lData
{-# INLINE outerLength #-}


instance Ragged L Ix1 e where
  flattenRagged = id
  {-# INLINE flattenRagged #-}
  generateRaggedM !comp !k f = do
    xs <-
      loopDeepM 0 (< coerce k) (+ 1) [] $ \i acc -> do
        e <- f i
        return (e : acc)
    return $ LArray comp $ coerce xs
  {-# INLINE generateRaggedM #-}
  loadRaggedST _scheduler xs uWrite start end sz = go (unList (lData xs)) start
    where
      go (y:ys) i
        | i < end = uWrite i y >> go ys (i + 1)
        | otherwise = throwM $ DimTooLongException 1 sz (outerLength xs)
      go [] i = when (i /= end) $ throwM $ DimTooShortException 1 sz (outerLength xs)
  {-# INLINE loadRaggedST #-}
  raggedFormat f _ arr = L.concat $ "[ " : L.intersperse ", " (map f (coerce (lData arr))) ++ [" ]"]


instance (Shape L ix, Ragged L ix e) => Load L ix e where
  makeArray comp sz f = runIdentity $ generateRaggedM comp sz (pure . f)
  {-# INLINE makeArray #-}
  iterArrayLinearST_ scheduler arr uWrite =
    loadRaggedST scheduler arr uWrite 0 (totalElem sz) sz
    where !sz = outerSize arr
  {-# INLINE iterArrayLinearST_ #-}

instance Ragged L Ix2 e where
  generateRaggedM = unsafeGenerateParM
  {-# INLINE generateRaggedM #-}
  flattenRagged arr = LArray {lComp = lComp arr, lData = coerce xs}
    where
      xs = concatMap (unList . lData . flattenRagged . LArray (lComp arr)) (unList (lData arr))
  {-# INLINE flattenRagged #-}
  loadRaggedST scheduler xs uWrite start end sz
    | isZeroSz sz = when (isNotNull (flattenRagged xs)) (throwM ShapeNonEmpty)
    | otherwise = do
      let (k, szL) = unconsSz sz
          step = totalElem szL
      leftOver <-
        loopM start (< end) (+ step) (coerce (lData xs)) $ \i zs ->
          case zs of
            [] -> throwM (DimTooShortException 2 k (outerLength xs))
            (y:ys) -> do
              scheduleWork_ scheduler $
                let end' = i + step
                    go (a:as) j
                      | j < end' = uWrite j a >> go as (j + 1)
                      | otherwise = throwM $ DimTooLongException 1 szL (Sz (length y))
                    go [] j = when (j /= end') $ throwM (DimTooShortException 1 szL (Sz (length y)))
                 in go y i
              pure ys
      unless (null leftOver) $ throwM $ DimTooLongException 2 k (outerLength xs)
  {-# INLINE loadRaggedST #-}
  raggedFormat f sep (LArray comp xs) =
    showN (\s y -> raggedFormat f s (LArray comp y :: Array L Ix1 e)) sep (coerce xs)

instance ( Shape L (IxN n)
         , Ragged L (Ix (n - 1)) e
         , Coercible (Elt (Ix (n - 1)) e) (ListItem (Ix (n - 1)) e)
         ) =>
         Ragged L (IxN n) e where
  generateRaggedM = unsafeGenerateParM
  {-# INLINE generateRaggedM #-}
  flattenRagged arr = LArray {lComp = lComp arr, lData = coerce xs}
    where
      xs = concatMap (unList . lData . flattenRagged . LArray (lComp arr)) (unList (lData arr))
  {-# INLINE flattenRagged #-}
  loadRaggedST scheduler xs uWrite start end sz
    | isZeroSz sz = when (isNotNull (flattenRagged xs)) (throwM ShapeNonEmpty)
    | otherwise = do
      let (k, szL) = unconsSz sz
          step = totalElem szL
          subScheduler
            | end - start < numWorkers scheduler * step = scheduler
            | otherwise = trivialScheduler_
      leftOver <-
        loopM start (< end) (+ step) (unList (lData xs)) $ \i zs ->
          case zs of
            [] -> throwM (DimTooShortException (dimensions sz) k (outerLength xs))
            (y:ys) -> do
              scheduleWork_ scheduler $
                loadRaggedST subScheduler (LArray Seq y) uWrite i (i + step) szL
              pure ys
      unless (null leftOver) $ throwM $ DimTooLongException (dimensions sz) k (outerLength xs)
  {-# INLINE loadRaggedST #-}
  raggedFormat f sep (LArray comp xs) =
    showN (\s y -> raggedFormat f s (LArray comp y :: Array L (Ix (n - 1)) e)) sep (coerce xs)

unsafeGenerateParM ::
     (Elt ix e ~ List (Lower ix) e, Index ix, Monad m, Ragged L (Lower ix) e)
  => Comp
  -> Sz ix
  -> (ix -> m e)
  -> m (Array L ix e)
unsafeGenerateParM comp !sz f = do
  res <- sequence $ unsafePerformIO $ do
    let !(ksz, szL) = unconsSz sz
        !k = unSz ksz
    withScheduler comp $ \ scheduler ->
      splitLinearly (numWorkers scheduler) k $ \ chunkLength slackStart -> do
        loopA_ 0 (< slackStart) (+ chunkLength) $ \ !start ->
          scheduleWork scheduler $ do
            res <- loopDeepM start (< (start + chunkLength)) (+ 1) [] $ \i acc ->
              return (fmap lData (generateRaggedM Seq szL (\ !ixL -> f (consDim i ixL))):acc)
            return $! sequence res
        when (slackStart < k) $
          scheduleWork scheduler $ do
            res <- loopDeepM slackStart (< k) (+ 1) [] $ \i acc ->
              return (fmap lData (generateRaggedM Seq szL (\ !ixL -> f (consDim i ixL))):acc)
            return $! sequence res
  return $ LArray comp $ List $ concat res
{-# INLINE unsafeGenerateParM #-}

instance Strategy L where
  setComp c arr = arr {lComp = c}
  {-# INLINE setComp #-}
  getComp = lComp
  {-# INLINE getComp #-}
  repr = L

-- -- TODO: benchmark against using unsafeGenerateM directly
-- unsafeGenerateN ::
--   ( Ragged r ix e
--   , Ragged r (Lower ix) e
--   , Elt r ix e ~ Array r (Lower ix) e )
--   => Comp
--   -> Sz ix
--   -> (ix -> e)
--   -> Array r ix e
-- unsafeGenerateN comp sz f = unsafePerformIO $ do
--   let !(m, szL) = unconsSz sz
--   xs <- withScheduler comp $ \scheduler ->
--     loopM_ 0 (< coerce m) (+ 1) $ \i -> scheduleWork scheduler $
--       generateRaggedM comp szL $ \ix -> return $ f (consDim i ix)
--   return $! foldr' consR (emptyR comp) xs
-- {-# INLINE unsafeGenerateN #-}


-- | Construct an array backed by linked lists from any source array
--
-- @since 0.4.0
toListArray :: (Ragged L ix e, Shape r ix, Source r e) => Array r ix e -> Array L ix e
toListArray !arr = makeArray (getComp arr) (outerSize arr) (unsafeIndex arr)
{-# INLINE toListArray #-}



instance (Ragged L ix e, Show e) => Show (Array L ix e) where
  showsPrec n arr  = showsArrayLAsPrec (Proxy :: Proxy L) (outerSize arr) n arr

instance (Ragged L ix e, Show e) => Show (List ix e) where
  show xs = "  " ++ raggedFormat show "\n  " arrL
    where arrL = LArray Seq xs :: Array L ix e


showN :: (String -> a -> String) -> String -> [a] -> String
showN _     _        [] = "[  ]"
showN fShow lnPrefix ls =
  L.concat
    (["[ "] ++
     L.intersperse (lnPrefix ++ ", ") (map (fShow (lnPrefix ++ "  ")) ls) ++ [lnPrefix, "]"])


showsArrayLAsPrec ::
     forall r ix e. (Ragged L ix e, Typeable r, Show e)
  => Proxy r
  -> Sz ix
  -> Int
  -> Array L ix e -- Array to show
  -> ShowS
showsArrayLAsPrec pr sz n arr =
  opp .
  ("Array " ++) .
  showsTypeRep (typeRep pr) .
  (' ':) .
  showsPrec 1 (getComp arr) . (" (" ++) . shows sz . (")\n" ++) . shows lnarr . clp
  where
    (opp, clp) =
      if n == 0
        then (id, id)
        else (('(':), ("\n)" ++))
    lnarr = lData arr

-- | Helper function for declaring `Show` instances for arrays
--
-- @since 0.4.0
showsArrayPrec ::
     forall r r' ix e. (Ragged L ix e, Load r ix e, Load r' ix e, Source r' e, Show e)
  => (Array r ix e -> Array r' ix e) -- ^ Modifier
  -> Int
  -> Array r ix e -- Array to show
  -> ShowS
showsArrayPrec f n arr = showsArrayLAsPrec (Proxy :: Proxy r) sz n larr
  where
    sz = size arr'
    arr' = f arr
    larr = makeArray (getComp arr') sz (evaluate' arr') :: Array L ix e


-- | Helper function for declaring `Show` instances for arrays
--
-- @since 0.4.0
showArrayList
  :: Show arr => [arr] -> String -> String
showArrayList arrs = ('[':) . go arrs . (']':)
  where
    go []     = id
    go [x]    = (' ':) . shows x . ('\n':)
    go (x:xs) = (' ':) . shows x . ("\n," ++) . go xs


instance Stream L Ix1 e where
  toStream = S.fromList . unList . lData
  {-# INLINE toStream #-}
  toStreamIx = S.indexed . S.fromList . unList . lData
  {-# INLINE toStreamIx #-}