content: {-# OPTIONS_HADDOCK not-home #-}
  -- |
  -- Copyright: (c) 2021 Xy Ren
  -- License: BSD3
  -- Maintainer: xy.r@outlook.com
  -- Stability: unstable
  -- Portability: non-portable (GHC only)
  module Data.RadixVec (RadixVec, size, empty, lookup, update, snoc) where
  
  import           Control.Monad.ST          (ST)
  import           Data.Bits                 (Bits (unsafeShiftL, unsafeShiftR, (.&.)), FiniteBits (countTrailingZeros))
  import           Data.Primitive.MachDeps   (sIZEOF_INT)
  import           Data.Primitive.SmallArray (SmallArray, SmallMutableArray, copySmallArray, indexSmallArray,
                                              newSmallArray, readSmallArray, runSmallArray, sizeofSmallArray,
                                              thawSmallArray, writeSmallArray)
  import           Prelude                   hiding (lookup)
  
  -- | An efficient vector type, implemented as a radix tree. It has the following time complexities:
  --
  -- * Lookup: \( O(\log n) \)
  -- * Update: \( O(\log n) \)
  -- * Append: \( O(\log n) \)
  --
  -- The branching factor (base of log) is 32 therefore the time is close to constant in most cases. Note that in
  -- practice, lookup is faster than update, and update is faster than append.
  data RadixVec a = RadixVec !Int !(Tree a)
  
  type Shift = Int
  
  -- | The \( \log_2 \) of the branching factor. The branching factor is set to be 32 for now but may change in the
  -- future.
  factor :: Int
  factor = 5
  
  -- | A mask covering one chunk of an index.
  initialMask :: Int
  initialMask = (1 `unsafeShiftL` factor) - 1
  
  -- | A radix tree. The tree is always left-leaning.
  data Tree a
    = Tip
      {-# UNPACK #-} !(SmallArray a)
    | Node
      {-# UNPACK #-} !Shift
      {-# UNPACK #-} !(SmallArray (Tree a))
  
  -- | Mask a portion of an index.
  mask :: Shift -> Int -> Int
  mask s x = initialMask .&. (x `unsafeShiftR` s)
  
  -- | Mask the zeroth portion of the index.
  mask0 :: Int -> Int
  mask0 x = initialMask .&. x
  
  -- | Alter an element in a 'SmallMutableArray' by a function.
  alterSmallArray :: SmallMutableArray s a -> Int -> (a -> a) -> ST s ()
  alterSmallArray marr ix f = do
    x <- readSmallArray marr ix
    writeSmallArray marr ix $! f x
  
  -- | Get the size of the 'RadixVec'.
  size :: RadixVec a -> Int
  size (RadixVec sz _) = sz
  
  -- | The empty 'RadixVec'.
  empty :: RadixVec a
  empty = RadixVec 0 $ Tip $ runSmallArray $ newSmallArray 0 $ error
    "Cleff.Internal.RadixVec: Encountered an element in an empty RadixVec. Please report this as a bug."
  
  -- | Lookup in a 'RadixVec' by an index. This does not perform any bounds check.
  lookup :: Int -> RadixVec a -> a
  lookup ix (RadixVec _ tree) = go tree
    where
      go (Tip arr)    = indexSmallArray arr (initialMask .&. ix)
      go (Node s arr) = go (indexSmallArray arr (mask s ix))
  
  -- | Update a value in a 'RadixVec' by an index. The value will be forced before installing. This does not perform any
  -- bounds check.
  update :: Int -> a -> RadixVec a -> RadixVec a
  update ix x (RadixVec len tree) = RadixVec len (go tree)
    where
      go (Tip arr) = Tip $ runSmallArray do
        marr <- thawSmallArray arr 0 (sizeofSmallArray arr)
        writeSmallArray marr (mask0 ix) $! x
        pure marr
      go (Node s arr) = Node s $ runSmallArray do
        marr <- thawSmallArray arr 0 (sizeofSmallArray arr)
        alterSmallArray marr (mask s ix) go
        pure marr
  
  -- | Append a value to a 'RadixVec'. The value will be forced before installing. This does not perform any bounds check.
  snoc :: RadixVec a -> a -> RadixVec a
  snoc (RadixVec len tree) x
    | ins <= topShift = RadixVec (len + 1) (go tree)
    | otherwise = RadixVec (len + 1) $ Node (topShift + factor) $ runSmallArray $ do
      marr <- newSmallArray 2 $! tree
      writeSmallArray marr 1 $! branch topShift
      pure marr
    where
      topShift = case tree of
        Tip _    -> 0
        Node s _ -> s
      ins = (countTrailingZeros len `mod` sIZEOF_INT `div` factor) * factor
      branch 0 = Tip $ runSmallArray $ newSmallArray 1 $! x
      branch s = Node s $ runSmallArray $ newSmallArray 1 $! branch (s - factor)
      enlarge arr new = runSmallArray do
        let sz = sizeofSmallArray arr
        marr <- newSmallArray (sz + 1) $! new
        copySmallArray marr 0 arr 0 sz
        pure marr
      go (Tip arr) = Tip $ enlarge arr x
      go (Node s arr)
        | ins == s = Node s $ enlarge arr $ branch (s - factor)
        | otherwise = Node s $ runSmallArray do
          marr <- thawSmallArray arr 0 (sizeofSmallArray arr)
          alterSmallArray marr (mask s len) go
          pure marr
  