-- | Persistent array using SmallArray# nodes and a tail cache, in the style
--   of Clojure persistent arrays.
--
--   The branching factor W is chosen with CPP: compile with -DBRANCHING=4, 8 or
--   16. The default is 8.
module Utils.TreeArray (
    Array (..)
  , empty
  , index
  , modify
  , size
  , snoc
  , update
  , same
  ) where

#ifndef BRANCHING
#define BRANCHING 8
#endif

import GHC.Exts hiding (toList)
import Data.Bits
import Data.Kind (Type)

--------------------------------------------------------------------------------
-- Branching-dependent definitions

-- | log2 of the branching factor.
bits :: Int

-- | Shift of the root node for a tree of the given (positive) number of
--   chunks: the bit length of the largest chunk index (m - 1), rounded up to a
--   multiple of "bits", minus "bits". A 1-chunk tree has shift (-bits): its
--   root is the chunk itself.
rootShift :: Int -> Int

-- | Every "bits"-th bit set: m is a power of W iff it is a power of two and
--   has one of these bits set (this also rules out m == 0).
powMask :: Int

#if BRANCHING == 4
bits = 2
rootShift m = ((65 - countLeadingZeros (m - 1)) .&. (-2)) - 2
powMask = 0x5555555555555555
#elif BRANCHING == 8
bits = 3
-- x `quot` 3 == (x * 43) >> 7 for 0 <= x < 128, and here 2 <= x <= 65;
-- q * 3 == q + q << 1. GHC's native code generator would emit idiv and imul.
rootShift m =
  let q = unsafeShiftR ((66 - countLeadingZeros (m - 1)) * 43) 7
  in q + unsafeShiftL q 1 - 3
powMask = 0x1249249249249249
#elif BRANCHING == 16
bits = 4
rootShift m = ((67 - countLeadingZeros (m - 1)) .&. (-4)) - 4
powMask = 0x1111111111111111
#else
#error "BRANCHING must be 4, 8 or 16"
#endif

{-# inline bits #-}
{-# inline rootShift #-}
{-# inline powMask #-}

{-# inline width #-}
width :: Int
width = unsafeShiftL 1 bits

{-# inline mask #-}
mask :: Int
mask = width - 1


--------------------------------------------------------------------------------
-- Nodes

type U = (Any :: UnliftedType)
type Node = SmallArray# U

type role Array representational
-- | Chunk tree, tail, size.
data Array (a :: Type) = Array U Node Int

{-# inline toU #-}
toU :: a -> U
toU x = unsafeCoerce# x -- x is forced by Strict

{-# inline fromU #-}
fromU :: U -> a
fromU x = unsafeCoerce# x

{-# inline nodeU #-}
nodeU :: Node -> U
nodeU n = unsafeCoerce# n

{-# inline uNode #-}
uNode :: U -> Node
uNode t = unsafeCoerce# t

-- | Filler for unused slots.
{-# inline hole #-}
hole :: (# #) -> U
hole _ = toU ()

{-# inline get #-}
get :: Node -> Int -> U
get n (I# i) = case indexSmallArray# n i of (# x #) -> x

{-# inline set #-}
-- | Copy of the node with one child replaced.
set :: Node -> Int -> U -> Node
set n (I# i) v = case width of
  I# w -> runRW# \s -> case thawSmallArray# n 0# w s of
    (# s, m #) -> case writeSmallArray# m i v s of
      s -> case unsafeFreezeSmallArray# m s of
        (# _, n' #) -> n'

{-# inline node1 #-}
-- | Node with the given first child and holes elsewhere.
node1 :: U -> Node
node1 a = case width of
  I# w -> runRW# \s -> case newSmallArray# w (hole (# #)) s of
    (# s, m #) -> case writeSmallArray# m 0# a s of
      s -> case unsafeFreezeSmallArray# m s of
        (# _, n #) -> n

{-# inline node2 #-}
-- | Node with the given first two children and holes elsewhere.
node2 :: U -> U -> Node
node2 a b = case width of
  I# w -> runRW# \s -> case newSmallArray# w (hole (# #)) s of
    (# s, m #) -> case writeSmallArray# m 0# a s of
      s -> case writeSmallArray# m 1# b s of
        s -> case unsafeFreezeSmallArray# m s of
          (# _, n #) -> n

--------------------------------------------------------------------------------
-- Operations

empty :: Array a
empty = Array (hole (# #)) (node1 (hole (# #))) 0

-- | 0 <= x < s, as a single unsigned comparison.
inBounds :: Int -> Int -> Bool
inBounds x s = (fromIntegral x :: Word) < fromIntegral s

{-# inline digit #-}
digit :: Int -> Int -> Int
digit x shift = unsafeShiftR x shift .&. mask

{-# inline size #-}
size :: Array a -> Int
size (Array _ _ s) = s

{-# inline same #-}
same :: Array a -> Array a -> Bool
same (Array t tl _) (Array t' tl' _) =
     isTrue# (reallyUnsafePtrEquality# t t')
  && isTrue# (reallyUnsafePtrEquality# tl tl')

{-# inline treeSize #-}
-- | Number of elements stored in the tree.
treeSize :: Int -> Int
treeSize s = s .&. complement mask

-- | Chunk number c of a tree with root at "shift".
chunk :: U -> Int -> Int -> Node
chunk t shift c | shift < 0 = uNode t
chunk t shift c = chunk (get (uNode t) (digit c shift)) (shift - bits) c

index :: Array a -> Int -> a
index (Array t tl s) x
  | inBounds x s =
    if x >= treeSize s then fromU (get tl (x .&. mask))
    else fromU (get (chunk t (rootShift (unsafeShiftR s bits)) (unsafeShiftR x bits)) (x .&. mask))
index _ _ = error "TreeArray.index: out of bounds"

update :: Array a -> Int -> a -> Array a
update (Array t tl s) x new
  | inBounds x s =
    if x >= treeSize s then Array t (set tl (x .&. mask) (toU new)) s
    else Array (updateGo (unsafeShiftR x bits) (x .&. mask) (toU new) t (rootShift (unsafeShiftR s bits))) tl s
update _ _ _ = error "TreeArray.update: out of bounds"

-- | Set element j of chunk c.
updateGo :: Int -> Int -> U -> U -> Int -> U
updateGo c j new t shift | shift < 0 = nodeU (set (uNode t) j new)
updateGo c j new t shift =
  let n = uNode t; i = digit c shift in
  nodeU (set n i (updateGo c j new (get n i) (shift - bits)))

{-# inline modify #-}
modify :: forall a. Array a -> Int -> (a -> a) -> Array a
modify (Array t tl s) x f
  | inBounds x s =
    if x >= treeSize s then Array t (leaf tl) s
    else Array (go t (rootShift (unsafeShiftR s bits))) tl s
  where
    j = x .&. mask
    c = unsafeShiftR x bits
    -- "f" stays captured, so that it is inlined here when modify is
    leaf :: Node -> Node
    leaf n = let x = fromU (get n j) in set n j (toU (f x))
    go :: U -> Int -> U
    go t shift | shift < 0 = nodeU (leaf (uNode t))
    go t shift =
      let n = uNode t; i = digit c shift in
      nodeU (set n i (go (get n i) (shift - bits)))
modify _ _ _ = error "TreeArray.modify: out of bounds"

snoc :: Array a -> a -> Array a
snoc (Array t tl s) new =
  let j   = s .&. mask
      tl' = set tl j (toU new)
  in if j /= mask then Array t tl' (s + 1)
     -- the tail is full: push it into the tree. tl' is also reused as the new
     -- (logically empty) tail, since its slots are overwritten before being read.
     else Array (pushChunk t (unsafeShiftR s bits) (nodeU tl')) tl' (s + 1)

-- | Append a chunk to a tree of m chunks. The tree is full exactly when m is a
--   power of W (including 1).
pushChunk :: U -> Int -> U -> U
pushChunk t m new
  | m .&. (m - 1) == 0 && m .&. powMask /= 0 =
    nodeU (node2 t (snocPath new (rootShift m)))
  | otherwise =
    snocGo m (countTrailingZeros m - bits) new t (rootShift (m + 1))

-- | A fresh path of nodes leading to "new" at index 0.
snocPath :: U -> Int -> U
snocPath new shift | shift < 0 = new
snocPath new shift = nodeU (node1 (snocPath new (shift - bits)))

-- | Insert "new" at index m, given the shift at and below which subtrees are
--   fresh (the low (shift + bits) bits of m are zero).
snocGo :: Int -> Int -> U -> U -> Int -> U
snocGo m fresh new t shift | shift <= fresh = snocPath new shift
snocGo m fresh new t shift =
  let n = uNode t; i = digit m shift in
  nodeU (set n i (snocGo m fresh new (get n i) (shift - bits)))

{-# inline foldr' #-}
foldr' :: forall a b. (a -> b -> b) -> b -> Array a -> b
foldr' = \f b (Array t tl s) ->

  let -- elements i, i-1 .. 0 of a node
      elems :: Node -> Int -> b -> b
      elems n i b | i < 0 = b
      elems n i b = let x = fromU (get n i) in elems n (i - 1) (f x b)

      -- children i, i-1 .. 0 of a node
      fulls :: Node -> Int -> Int -> b -> b
      fulls n shift i b | i < 0 = b
      fulls n shift i b = fulls n shift (i - 1) (full (get n i) shift b)

      notfull :: U -> Int -> Int -> b -> b
      notfull t shift m b | shift < 0 = elems (uNode t) mask b
      notfull t shift m b =
        let n = uNode t; shift' = shift - bits; ix = digit m shift in
        fulls n shift' (ix - 1) (notfull (get n ix) shift' m b)

      full :: U -> Int -> b -> b
      full t shift b | shift < 0 = elems (uNode t) mask b
      full t shift b = fulls (uNode t) (shift - bits) mask b

      m  = unsafeShiftR s bits
      b' = elems tl ((s .&. mask) - 1) b

  in case m of
    0 -> b'
    _ -> notfull t (rootShift m) (m - 1) b'

toList :: Array a -> [a]
toList = foldr' (:) []

instance Show a => Show (Array a) where
  show = show . toList