Fix wrong insertShift computation in <|

src/Data/RRBVector/Internal.hs

@@ -174,7 +174,7 @@
           where
             subtree = A.index arr i
 
--- Integer log base 2.
+-- | Integer log base 2.
 log2 :: Int -> Int
 log2 x = bitSizeMinus1 - countLeadingZeros x
   where
@@ -716,10 +716,10 @@
 
         viewrArr arr = (A.take arr (length arr - 1), A.last arr)
 
     -- the type signature is necessary to compile
     mergeRebalance :: forall a. Shift -> A.Array (Tree a) -> A.Array (Tree a) -> A.Array (Tree a) -> A.Array (Tree a)
     mergeRebalance !sh !left !center !right
         | sh == blockShift = mergeRebalance' (\(Leaf arr) -> arr) Leaf
         | otherwise = mergeRebalance' treeToArray (computeSizes (down sh))
       where
         mergeRebalance' :: (Tree a -> A.Array t) -> (A.Array t -> Tree a) -> A.Array (Tree a)
@@ -768,7 +768,10 @@
     -- compute the shift at which the new branch needs to be inserted (0 means there is space in the leaf)
     -- the size is computed for efficient calculation of the shift in a balanced subtree
     computeShift !sz !sh !min (Balanced _) =
-        let newShift = (log2 sz `div` blockShift) * blockShift
+        -- @sz - 1@ is the index of the last element
+        let hiShift = max ((log2 (sz - 1) `div` blockShift) * blockShift) 0 -- the shift of the root when normalizing
+            hi = (sz - 1) `unsafeShiftR` hiShift -- the length of the root node when normalizing minus 1
+            newShift = if hi < blockMask then hiShift else hiShift + blockShift
         in if newShift > sh then min else newShift
     computeShift _ sh min (Unbalanced arr sizes) =
         let sz' = indexPrimArray sizes 0 -- the size of the first subtree