Ix.hs

ix :: Applicative f => Int -> (a -> f a) -> Seq a -> f (Seq a)
ix i@(I# i') f (Seq xs)
  | 0 <= i && i < size xs = Seq <$> ixTreeE (\_ (Elem a) -> Elem <$> f a) i' xs
  | otherwise   = pure (Seq xs)

unInt :: Int -> Int#
unInt (I# n) = n

ixTreeE :: Applicative f
       => (Int# -> Elem a -> f (Elem a)) -> Int# -> FingerTree (Elem a) -> f (FingerTree (Elem a))
ixTreeE _ _ EmptyT = pure EmptyT
ixTreeE f i (Single x) = Single <$> f i x
ixTreeE f i (Deep s pr m sf)
  | I# i < spr  = (\q -> Deep s q m sf) <$> ixDigitE f i pr
  | I# i < spm  = (\q -> Deep s pr q sf) <$> ixTreeN (ixNodeE f) (unInt $ I# i - spr) m
  | otherwise   = (\q -> Deep s pr m q) <$> ixDigitE f (unInt $ I# i - spm) sf
  where
    spr     = size pr
    spm     = spr + size m

ixTreeN :: Applicative f
       => (Int# -> Node a -> f (Node a)) -> Int# -> FingerTree (Node a) -> f (FingerTree (Node a))
ixTreeN _ _ EmptyT = pure EmptyT
ixTreeN f i (Single x) = Single <$> f i x
ixTreeN f i (Deep s pr m sf)
  | I# i < spr  = (\q -> Deep s q m sf) <$> ixDigitN f i pr
  | I# i < spm  = (\q -> Deep s pr q sf) <$> ixTreeN (ixNodeN f) (unInt $ I# i - spr) m
  | otherwise   = (\q -> Deep s pr m q) <$> ixDigitN f (unInt $ I# i - spm) sf
  where
    spr     = size pr
    spm     = spr + size m

ixNodeE :: Applicative f => (Int# -> Elem a -> f (Elem a)) -> Int# -> Node (Elem a) -> f (Node (Elem a))
ixNodeE f i t = ixNode f i t

ixNodeN :: Applicative f => (Int# -> Node a -> f (Node a)) -> Int# -> Node (Node a) -> f (Node (Node a))
ixNodeN f i t = ixNode f i t

{-# INLINE ixNode #-}
ixNode :: (Applicative f, Sized a) => (Int# -> a -> f a) -> Int# -> Node a -> f (Node a)
ixNode f i (Node2 s a b)
  | I# i < sa   = (\q -> Node2 s q b) <$> f i a
  | otherwise   = (\q -> Node2 s a q) <$> f (unInt $ I# i - sa) b
  where
    sa      = size a
ixNode f i (Node3 s a b c)
  | I# i < sa   = (\q -> Node3 s q b c) <$> f i a
  | I# i < sab  = (\q -> Node3 s a q c) <$> f (unInt $ I# i - sa) b
  | otherwise   = (\q -> Node3 s a b q) <$> f (unInt $ I# i - sab) c
  where
    sa      = size a
    sab     = sa + size b

ixDigitE :: Applicative f => (Int# -> Elem a -> f (Elem a)) -> Int# -> Digit (Elem a) -> f (Digit (Elem a))
ixDigitE f i t = ixDigit f i t
ixDigitN :: Applicative f => (Int# -> Node a -> f (Node a)) -> Int# -> Digit (Node a) -> f (Digit (Node a))
ixDigitN f i t = ixDigit f i t

{-# INLINE ixDigit #-}
ixDigit :: (Applicative f, Sized a) => (Int# -> a -> f a) -> Int# -> Digit a -> f (Digit a)
ixDigit f i (One a) = One <$> f i a
ixDigit f i (Two a b)
  | I# i < sa   = (\q -> Two q b) <$> f i a
  | otherwise   = (\q -> Two a q) <$> f (unInt $ I# i - sa) b
  where
    sa      = size a
ixDigit f i (Three a b c)
  | I# i < sa   = (\q -> Three q b c) <$> f i a
  | I# i < sab  = (\q -> Three a q c) <$> f (unInt $ I# i - sa) b
  | otherwise   = (\q -> Three a b q) <$> f (unInt $ I# i - sab) c
  where
    sa      = size a
    sab     = sa + size b
ixDigit f i (Four a b c d)
  | I# i < sa   = (\q -> Four q b c d) <$> f i a
  | I# i < sab  = (\q -> Four a q c d) <$> f (unInt $ I# i - sa) b
  | I# i < sabc = (\q -> Four a b q d) <$> f (unInt $ I# i - sab) c
  | otherwise   = (\q -> Four a b c q) <$> f (unInt $ I# i - sabc) d
  where
    sa      = size a
    sab     = sa + size b
    sabc    = sab + size c