module ImmutGrowVector where

import Data.Vector.Unboxed (Unbox)
import Data.Vector.Unboxed qualified as VU


-- | 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 (VU.Vector a) (VU.Vector a)
  deriving (Show)

empty :: Unbox a => ImmutGrowVector a
empty = ImmutGrowVector VU.empty VU.empty

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

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

length :: Unbox a => ImmutGrowVector a -> Int
length (ImmutGrowVector prefix suffix) = VU.length prefix + VU.length suffix

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

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

-- | Transform the vector by mapping the underlying unboxed vectors
mapUVector :: (VU.Vector a -> VU.Vector b) -> ImmutGrowVector a -> ImmutGrowVector b
mapUVector f (ImmutGrowVector prefix suffix) = ImmutGrowVector (f prefix) (f suffix)