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 (IsList (..))
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 #-}