module ImmutGrowVector where

import Foreign.Storable
import Data.Vector.Storable qualified as VS


-- | Acts like an immutable storable vector, except that it's split in a long
-- prefix and a short suffix so that modifications at the end are cheap. If the
-- suffix gets longer (by appending elements), the suffix elements are promoted
-- to prefix elements once in a while, resulting in a big copy at those times.
data ImmutGrowVector a = ImmutGrowVector (VS.Vector a) (VS.Vector a)
  deriving (Show)

empty :: Storable a => ImmutGrowVector a
empty = ImmutGrowVector VS.empty VS.empty

fromListN :: Storable a => Int -> [a] -> ImmutGrowVector a
fromListN n l
  | n > 2 =
      let (l1, l2) = splitAt (n - 2) l
      in ImmutGrowVector (VS.fromListN (n - 2) l1) (VS.fromListN 2 l2)
  | otherwise =
      ImmutGrowVector VS.empty (VS.fromListN n l)

(!) :: Storable a => ImmutGrowVector a -> Int -> a
ImmutGrowVector prefix suffix ! i
  | i < VS.length prefix = prefix VS.! i
  | otherwise = suffix VS.! (i - VS.length prefix)

length :: Storable a => ImmutGrowVector a -> Int
length (ImmutGrowVector prefix suffix) = VS.length prefix + VS.length suffix

set :: Storable a => ImmutGrowVector a -> Int -> a -> ImmutGrowVector a
set (ImmutGrowVector prefix suffix) idx value
  | idx < VS.length prefix = error "ImmutGrowVector: mutation in slow part"
  | otherwise = ImmutGrowVector prefix (suffix VS.// [(idx - VS.length prefix, value)])

append :: Storable a => ImmutGrowVector a -> a -> ImmutGrowVector a
append (ImmutGrowVector prefix suffix) value
  | VS.length suffix < 8 = ImmutGrowVector prefix (suffix `VS.snoc` value)
  | otherwise =
      let n = VS.length suffix
      in ImmutGrowVector (prefix <> VS.take (n - 1) suffix) (VS.drop (n - 1) suffix `VS.snoc` value)