content: {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE LinearTypes #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE NumericUnderscores #-}
  {-# LANGUAGE QuantifiedConstraints #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE StandaloneKindSignatures #-}
  
  -- Uncomment the line below to observe the generated (optimised) Core. It will
  -- land in a file named “Array.dump-simpl”
  -- {-# OPTIONS_GHC -ddump-simpl -ddump-to-file -dsuppress-all -dsuppress-uniques #-}
  
  module Data.Mutable.Array (benchmarks) where
  
  import Control.DeepSeq (rnf)
  import qualified Data.Array.Mutable.Linear as Array.Linear
  import qualified Data.Array.Mutable.Linear as Array.Linear.Array
  import qualified Data.Foldable
  import Data.Functor.Compose
  import Data.Kind
  import qualified Data.Sequence
  import qualified Data.Unrestricted.Linear as Linear
  import qualified Data.Vector
  import Prelude.Linear (($), (&))
  import qualified Prelude.Linear as Linear
  import Test.Tasty.Bench
  import Prelude hiding (($))
  
  arr_size :: Int
  arr_size = 1_000
  
  benchmarks :: Benchmark
  benchmarks =
    bgroup
      "arrays"
      $ runImpls
        [ bAlloc,
          bToList,
          bMap,
          bReads,
          bSets
        ]
  
  --------------------------------------------------------------------------------
  
  data Impl where
    Impl :: String -> (forall arr. (ArrayThing arr) => arr Int %1 -> ()) -> Impl
  
  runImpls :: [Impl] -> [Benchmark]
  runImpls = map (runImpl arr_size)
  
  runImpl :: Int -> Impl -> Benchmark
  runImpl sz0 (Impl name impl) =
    bgroup
      name
      [ bench "Data.Array.Mutable.Linear" $ whnf (runLinear impl) sz0,
        bench "Data.Vector" $ whnf (runDataVector (cleanup impl)) sz0,
        bench "Data.Sequence" $ whnf (runSequence (cleanup impl)) sz0
      ]
    where
      runLinear :: (Array.Linear.Array Int %1 -> ()) -> Int -> ()
      runLinear cb sz = Linear.unur (Array.Linear.alloc sz 0 (\a -> Linear.move (cb a)))
  
      runDataVector :: (Data.Vector.Vector Int -> ()) -> Int -> ()
      runDataVector cb sz = cb (Data.Vector.replicate sz 0)
  
      runSequence :: (Data.Sequence.Seq Int -> ()) -> Int -> ()
      runSequence cb sz = cb (Data.Sequence.replicate sz 0)
  {-# INLINE runImpl #-}
  
  type ArrayThing :: (Type -> Type) -> Constraint
  class ArrayThing arr where
    size :: arr a %1 -> (Linear.Ur Int, arr a)
    get :: Int -> arr a %1 -> (Linear.Ur a, arr a)
    set :: Int -> a -> arr a %1 -> arr a
    toList :: arr a %1 -> Linear.Ur [a]
    amap :: (a -> b) -> arr a %1 -> arr b
  
    -- | Note: I [Arnaud Spiwack] initially thought I could use
    -- 'Consumable'/'consume' for this. But it doesn't work because the natural
    -- 'consume' function for `Ur x` doesn't evaluate the `x` at all. We need to
    -- evaluate the `x` in the 'Vector' instance.
    force :: arr a %1 -> ()
  
  type UArrayThing :: (Type -> Type) -> Constraint
  class UArrayThing arr where
    usize :: arr a -> Int
    uget :: Int -> arr a -> a
    uset :: Int -> a -> arr a -> arr a
    utoList :: arr a -> [a]
    uamap :: (a -> b) -> arr a -> arr b
    uforce :: arr a -> ()
  
  instance ArrayThing Array.Linear.Array where
    size = Array.Linear.Array.size
    get = Array.Linear.Array.unsafeGet
    set = Array.Linear.Array.unsafeSet
    toList = Array.Linear.Array.toList
    amap = Array.Linear.Array.map
    force = Linear.consume
  
  instance (UArrayThing arr) => ArrayThing (Compose Linear.Ur arr) where
    size (Compose (Linear.Ur arr)) = (Linear.Ur (usize arr), Compose (Linear.Ur arr))
    get i (Compose (Linear.Ur arr)) = (Linear.Ur (uget i arr), Compose (Linear.Ur arr))
    set i a (Compose (Linear.Ur arr)) = Compose (Linear.Ur (uset i a arr))
    toList (Compose (Linear.Ur arr)) = Linear.Ur (utoList arr)
    amap f (Compose (Linear.Ur arr)) = Compose (Linear.Ur (uamap f arr))
    force (Compose (Linear.Ur arr)) = uforce arr
  
  instance UArrayThing Data.Vector.Vector where
    usize = Data.Vector.length
    uget i v = v Data.Vector.! i
    uset i a v = v Data.Vector.// [(i, a)]
    utoList = Data.Vector.toList
    uamap = Data.Vector.map
    uforce = (`seq` ())
  
  instance UArrayThing Data.Sequence.Seq where
    usize = Data.Sequence.length
    uget i s = Data.Sequence.index s i
    uset = Data.Sequence.update
    utoList = Data.Foldable.toList
    uamap = fmap
  
    -- I'm not sure about this one: on the one hand it forces the data structure
    -- to be allocated. On the other hand, it will do an extra traversal. Maybe
    -- there's a better comparison that can be done.
    uforce s = (foldMap (\_ -> Strict) s) `seq` ()
  
  cleanup :: ((Compose Linear.Ur f a) %1 -> b) -> (f a -> b)
  cleanup k a = k (Compose (Linear.Ur a))
  
  data Strict = Strict
  
  instance Semigroup Strict where
    Strict <> x = x
  
  instance Monoid Strict where
    mempty = Strict
  
  --------------------------------------------------------------------------------
  
  bToList :: Impl
  bToList = Impl "toList" impl
    where
      impl :: (ArrayThing arr) => arr Int %1 -> ()
      impl arr = arr & toList & Linear.lift rnf & Linear.unur
  
  bMap :: Impl
  bMap = Impl "map" impl
    where
      impl :: (ArrayThing arr) => arr Int %1 -> ()
      impl arr =
        case arr & amap (+ 1) & get 5 of
          (Linear.Ur _, arr') -> force arr'
  
  bReads :: Impl
  bReads = Impl "reads" impl
    where
      impl :: (ArrayThing arr) => arr Int %1 -> ()
      impl arr0 =
        case size arr0 of
          (Linear.Ur sz, arr) -> go 0 sz arr
        where
          go :: (ArrayThing arr) => Int -> Int -> arr Int %1 -> ()
          go start end arr
            | start < end =
                case get start arr of
                  (Linear.Ur i, arr') -> i `Linear.seq` go (start + 1) end arr'
            | otherwise = force arr
  
  bAlloc :: Impl
  bAlloc = Impl "alloc" impl
    where
      impl :: (ArrayThing arr) => arr Int %1 -> ()
      impl = force
  
  bSets :: Impl
  bSets = Impl "successive writes (very unfair to vector)" impl
    where
      impl :: (ArrayThing arr) => arr Int %1 -> ()
      impl arr0 =
        case size arr0 of
          (Linear.Ur sz, arr) -> go 0 sz arr
        where
          go :: (ArrayThing arr) => Int -> Int -> arr Int %1 -> ()
          go start end arr
            | start < end =
                go (start + 1) end Linear.$ set start 42 arr
            | otherwise = force arr
  