{-|
== Mutable Fenwick tree, or binary indexed tree

A data structure supporting point updates and range queries, or the opposite.
See Fenwick.hs for a purely functional version. FenwickMut is multiple times faster when used with
unboxed arrays (see benchmarks).

Sources:

* Peter M. Fenwick, "A New Data Structure for Cumulative Frequency Tables", 1994
  https://citeseerx.ist.psu.edu/viewdoc/summary?doi=10.1.1.14.8917
* https://en.wikipedia.org/wiki/Fenwick_tree

Let n = r - l + 1 where (l, r) is the range of the Fenwick tree.
The complexities assume (<>) takes O(1) time.

-}

module FenwickMut
    ( FenwickMut
    , emptyFM
    , mappendFM
    , foldPrefixFM
    , foldRangeFM
    , mappendRangeFM
    ) where

import Control.Monad
import Data.Array.Base
import Data.Bits

import Misc ( Commutative, Group(..), modifyArray' )

type FenwickMut marr a = marr Int a

-- | Builds a Fenwick tree on range (l, r) where each element is mempty. O(n).
emptyFM :: (Monoid a, MArray marr a m) => (Int, Int) -> m (FenwickMut marr a)
emptyFM :: (Int, Int) -> m (FenwickMut marr a)
emptyFM = ((Int, Int) -> a -> m (FenwickMut marr a))
-> a -> (Int, Int) -> m (FenwickMut marr a)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Int, Int) -> a -> m (FenwickMut marr a)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray a
forall a. Monoid a => a
mempty

-- | mappends to the element at an index. O(log n).
mappendFM :: (Monoid a, MArray marr a m) => FenwickMut marr a -> Int -> a -> m ()
mappendFM :: FenwickMut marr a -> Int -> a -> m ()
mappendFM FenwickMut marr a
a Int
i a
x = do
    (Int
l,Int
r) <- FenwickMut marr a -> m (Int, Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
getBounds FenwickMut marr a
a
    (Int -> m ()) -> [Int] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\Int
j -> FenwickMut marr a -> Int -> (a -> a) -> m ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> (e -> e) -> m ()
modifyArray' FenwickMut marr a
a (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
x)) ([Int] -> m ()) -> [Int] -> m ()
forall a b. (a -> b) -> a -> b
$
        (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
rInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> Int -> [Int]
forall a. (a -> a) -> a -> [a]
iterate (\Int
j -> Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (-Int
j)) (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

-- | The result of folding the prefix upto the given index. Indices outside the tree range are allowed,
-- it is assumed elements there are mempty. O(log n).
foldPrefixFM :: (Monoid a, MArray marr a m) => FenwickMut marr a -> Int -> m a
foldPrefixFM :: FenwickMut marr a -> Int -> m a
foldPrefixFM FenwickMut marr a
a Int
i = do
    (Int
l,Int
r) <- FenwickMut marr a -> m (Int, Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
getBounds FenwickMut marr a
a
    (a -> Int -> m a) -> a -> [Int] -> m a
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\a
z Int
j -> (a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
z) (a -> a) -> m a -> m a
forall (m :: * -> *) a b. Monad m => (a -> b) -> m a -> m b
<$!> FenwickMut marr a -> Int -> m a
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray FenwickMut marr a
a (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)) a
forall a. Monoid a => a
mempty ([Int] -> m a) -> [Int] -> m a
forall a b. (a -> b) -> a -> b
$
        (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>Int
0) ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> Int -> [Int]
forall a. (a -> a) -> a -> [a]
iterate (\Int
j -> Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
j Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. (-Int
j)) (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
i Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

-- | Folds the elements in the range (l, r). O(log n).
foldRangeFM :: (Commutative a, Group a, MArray marr a m) => FenwickMut marr a -> Int -> Int -> m a
foldRangeFM :: FenwickMut marr a -> Int -> Int -> m a
foldRangeFM FenwickMut marr a
a Int
l Int
r = do
    a
lx <- FenwickMut marr a -> Int -> m a
forall a (marr :: * -> * -> *) (m :: * -> *).
(Monoid a, MArray marr a m) =>
FenwickMut marr a -> Int -> m a
foldPrefixFM FenwickMut marr a
a (Int
lInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
    a
rx <- FenwickMut marr a -> Int -> m a
forall a (marr :: * -> * -> *) (m :: * -> *).
(Monoid a, MArray marr a m) =>
FenwickMut marr a -> Int -> m a
foldPrefixFM FenwickMut marr a
a Int
r
    a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> m a) -> a -> m a
forall a b. (a -> b) -> a -> b
$! a -> a
forall a. Group a => a -> a
invert a
lx a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
rx

-- | mappends to all elements in the range (l, r). Can be used with foldPrefixFM for point queries.
-- O(log n).
mappendRangeFM :: (Commutative a, Group a, MArray marr a m)
               => FenwickMut marr a -> Int -> Int -> a -> m ()
mappendRangeFM :: FenwickMut marr a -> Int -> Int -> a -> m ()
mappendRangeFM FenwickMut marr a
a Int
l Int
r a
x = do
    (Int
_,Int
r0) <- FenwickMut marr a -> m (Int, Int)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m (i, i)
getBounds FenwickMut marr a
a
    FenwickMut marr a -> Int -> a -> m ()
forall a (marr :: * -> * -> *) (m :: * -> *).
(Monoid a, MArray marr a m) =>
FenwickMut marr a -> Int -> a -> m ()
mappendFM FenwickMut marr a
a Int
l a
x
    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
r Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
r0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ FenwickMut marr a -> Int -> a -> m ()
forall a (marr :: * -> * -> *) (m :: * -> *).
(Monoid a, MArray marr a m) =>
FenwickMut marr a -> Int -> a -> m ()
mappendFM FenwickMut marr a
a (Int
rInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) (a -> a
forall a. Group a => a -> a
invert a
x)

--------------------------------------------------------------------------------
-- For tests

-- Allows specialization across modules
{-# INLINABLE mappendFM #-}
{-# INLINABLE foldPrefixFM #-}