module Data.FenwickTree where

import Control.Monad
import Control.Monad.Primitive
import Data.Bits
import Data.Function
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Mutable as UM

newtype FenwickTree s a = FenwickTree {forall s a. FenwickTree s a -> MVector s a
getFenwickTree :: UM.MVector s a}

newFenwickTree ::
  (U.Unbox a, Monoid a, PrimMonad m) =>
  Int ->
  m (FenwickTree (PrimState m) a)
newFenwickTree :: forall a (m :: * -> *).
(Unbox a, Monoid a, PrimMonad m) =>
Int -> m (FenwickTree (PrimState m) a)
newFenwickTree Int
n = MVector (PrimState m) a -> FenwickTree (PrimState m) a
forall s a. MVector s a -> FenwickTree s a
FenwickTree (MVector (PrimState m) a -> FenwickTree (PrimState m) a)
-> m (MVector (PrimState m) a) -> m (FenwickTree (PrimState m) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> a -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
UM.replicate (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a
forall a. Monoid a => a
mempty
{-# INLINE newFenwickTree #-}

-- | /O(n)/
buildFenwickTree ::
  (U.Unbox a, Monoid a, PrimMonad m) =>
  U.Vector a ->
  m (FenwickTree (PrimState m) a)
buildFenwickTree :: forall a (m :: * -> *).
(Unbox a, Monoid a, PrimMonad m) =>
Vector a -> m (FenwickTree (PrimState m) a)
buildFenwickTree Vector a
vec = do
  let n :: Int
n = Vector a -> Int
forall a. Unbox a => Vector a -> Int
U.length Vector a
vec
  ft <- Int -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
UM.unsafeNew (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  UM.write ft 0 mempty
  U.unsafeCopy (UM.tail ft) vec
  flip fix 1 $ \Int -> m ()
loop !Int
i -> Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    let j :: Int
j = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (-Int
i))
    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
j Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
      fti <- MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector (PrimState m) a
ft Int
i
      UM.unsafeModify ft (<> fti) j
    Int -> m ()
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  return $ FenwickTree ft
{-# INLINE buildFenwickTree #-}

{- | mappend [0..k)

 /O(log n)/
-}
mappendTo ::
  (PrimMonad m, U.Unbox a, Monoid a) =>
  FenwickTree (PrimState m) a ->
  Int ->
  m a
mappendTo :: forall (m :: * -> *) a.
(PrimMonad m, Unbox a, Monoid a) =>
FenwickTree (PrimState m) a -> Int -> m a
mappendTo (FenwickTree MVector (PrimState m) a
ft) = a -> Int -> m a
forall {m :: * -> *}.
(PrimState m ~ PrimState m, PrimMonad m) =>
a -> Int -> m a
go a
forall a. Monoid a => a
mempty
  where
    go :: a -> Int -> m a
go !a
acc !Int
i
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = do
          xi <- MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
UM.unsafeRead MVector (PrimState m) a
MVector (PrimState m) a
ft Int
i
          go (acc <> xi) (i - (i .&. (-i)))
      | Bool
otherwise = a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
acc
{-# INLINE mappendTo #-}

-- | /O(log n)/
mappendAt ::
  (U.Unbox a, Semigroup a, PrimMonad m) =>
  FenwickTree (PrimState m) a ->
  Int ->
  a ->
  m ()
mappendAt :: forall a (m :: * -> *).
(Unbox a, Semigroup a, PrimMonad m) =>
FenwickTree (PrimState m) a -> Int -> a -> m ()
mappendAt (FenwickTree MVector (PrimState m) a
ft) Int
k a
v = (((Int -> m ()) -> Int -> m ()) -> Int -> m ())
-> Int -> ((Int -> m ()) -> Int -> m ()) -> m ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Int -> m ()) -> Int -> m ()) -> Int -> m ()
forall a. (a -> a) -> a
fix (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (((Int -> m ()) -> Int -> m ()) -> m ())
-> ((Int -> m ()) -> Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int -> m ()
loop !Int
i -> do
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    MVector (PrimState m) a -> (a -> a) -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> (a -> a) -> Int -> m ()
UM.unsafeModify MVector (PrimState m) a
ft (a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
v) Int
i
    Int -> m ()
loop (Int -> m ()) -> Int -> m ()
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (-Int
i))
  where
    !n :: Int
n = MVector (PrimState m) a -> Int
forall a s. Unbox a => MVector s a -> Int
UM.length MVector (PrimState m) a
ft
{-# INLINE mappendAt #-}