replicateEach

ReplicateEach.hs

-- `beforeSeq` implements `xs <* ys`
beforeSeq :: Seq a -> Seq b -> Seq a
beforeSeq xs ys = replicateEach (length ys) xs

-- a wide tree with two extra nodes to its left and right
data Wide a = Wide a (FingerTree a) a

-- replicate each element of a sequence a fixed number of times
replicateEach :: Int -> Seq a -> Seq a
replicateEach n xs
    | n < 0  = error "replicateEach takes a nonnegative integer argument"
    | n == 0 = empty
    | n == 1 = xs
replicateEach n xs0 = case viewl xs0 of
    EmptyL   -> empty
    l :< xs1 -> case viewr xs1 of
        EmptyR -> replicate n l
        Seq rg  :> r -> case repEachWide n (Wide (Elem l) rg (Elem r)) of
            Wide (Elem l') rg' (Elem r') -> l' <| (Seq rg' |> r')

repEachWide :: Int -> Wide (Elem a) -> Wide (Elem a)
repEachWide n0 t0
    | n0 == 2 = flatWide (mapMulWide n0 n2 t0)
    | n0 == 3 = flatWide (mapMulWide n0 n3 t0)
    | otherwise = case n0 `quotRem` 3 of
       (q, 0) -> flatWide $ go t0 (q - 2) liftElem 3 3 3
           where liftElem e = let !n3e = n3 e in (n3e, n3e, n3e)
       (q, 1) -> flatWide $ go t0 (q - 1) liftElem 2 3 2
           where liftElem e = let !n2e = n2 e in (n2e, n3 e, n2e)
       (q, _) -> flatWide $ go t0 (q - 1) liftElem 3 3 2
           where liftElem e = let !n3e = n3 e in (n3e, n3e, n2 e)
  where
    n2 e = Node2 2 e e
    n3 e = Node3 3 e e e

    -- invariant: `lift` produces nodes of fixed sizes sl, sm, sr
    -- such that sl + n * sm + sr = n0
    go :: Wide (Elem a) -> Int -> (Elem a -> (Node c, Node c, Node c))
       -> Int -> Int -> Int -> Wide (Node c)
    go t 0 lift !_ !_ !_ = flatWide $ mapMulWide n0 lift' t
        where lift' e = let !(l, _, r) = lift e in Node2 n0 l r
    go t 1 lift !_ !_ !_  = flatWide $ mapMulWide n0 lift' t
        where lift' e = let !(l, m, r) = lift e in Node3 n0 l m r
    go t n lift sl sm sr = flatWide $ case n `quotRem` 3 of
        (q, 0) -> go t (q - 1) lift0 sl0 sm0 sr0
        (q, 1) -> go t (q - 1) lift1 sl1 sm1 sr1
        (q, _) -> go t q       lift2 sl2 sm2 sr2
      where
        (!sl0, !sm0, !sr0) = (sl + sm, 3*sm, 2*sm + sr)
        lift0 e = let !(l, m, r) = lift e in
                  (Node2 sl0 l m, Node3 sm0 m m m, Node3 sr0 m m r)
        (!sl1, !sm1, !sr1) = (sl + 2*sm, 3*sm, 2*sm + sr)
        lift1 e = let !(l, m, r) = lift e in
                  (Node3 sl1 l m m, Node3 sm1 m m m, Node3 sr1 m m r)
        (!sl2, !sm2, !sr2) = (sl + sm, 3*sm, sm + sr)
        lift2 e = let !(l, m, r) = lift e in
                  (Node2 sl2 l m, Node3 sm2 m m m, Node2 sr2 m r)

-- flatten one level of a wide finger tree
flatWide :: Sized a => Wide (Node a) -> Wide a
flatWide (Wide (Node2 _ l1 l2) ft (Node2 _ r1 r2)) =
    Wide l1 (deep (One l2) ft (One r1)) r2
flatWide (Wide (Node3 _ l1 l2 l3) ft (Node2 _ r1 r2)) =
    Wide l1 (deep (Two l2 l3) ft (One r1)) r2
flatWide (Wide (Node2 _ l1 l2) ft (Node3 _ r1 r2 r3)) =
    Wide l1 (deep (One l2) ft (Two r1 r2)) r3
flatWide (Wide (Node3 _ l1 l2 l3) ft (Node3 _ r1 r2 r3)) =
    Wide l1 (deep (Two l2 l3) ft (Two r1 r2)) r3
flatWide _ = error "flatWide: Not a rigid finger tree."
{-# INLINE flatWide #-}

-- mapMul variant for wide finger trees
mapMulWide :: Int -> (a -> b) -> Wide a -> Wide b
mapMulWide n f (Wide l ft r) = Wide (f l) (mapMulFT n f ft) (f r)
{-# INLINE mapMulWide #-}

ReplicateEachOpt.hs

-- `beforeSeq` implements `xs <* ys`
beforeSeq :: Seq a -> Seq b -> Seq a
beforeSeq xs ys = replicateEach (length ys) xs

-- a finger tree with two elements split off
data RepTop a = RepTop a (FingerTree a) a

-- replicate each element of a sequence a fixed number of times
replicateEach :: Int -> Seq a -> Seq a
replicateEach n xs
    | n < 0  = error "replicateEach takes a nonnegative integer argument"
    | n == 0 = empty
    | n == 1 = xs
replicateEach n xs0 = case viewl xs0 of
    EmptyL   -> empty
    l :< xs1 -> case viewr xs1 of
        EmptyR -> replicate n l
        Seq ft  :> r -> Seq (repEachTop n (RepTop (Elem l) ft (Elem r)))

-- bookkeeping for `repEachMiddle`
data RepMid a c = RepMid (Node c) (Node c) (FingerTree a) (Node c) (Node c)

repEachTop :: Int -> RepTop (Elem a) -> FingerTree (Elem a)
repEachTop n0 t0@(RepTop _ ft0 _)
    | n0 == 2 = flatTop (mapMulTop n0 n2 t0)
    | n0 == 3 = flatTop (mapMulTop n0 n3 t0)
    | otherwise = case n0 `quotRem` 3 of
        (q, 0) -> flatTop $ go (q - 2) t0 liftElem {- 3 3 3 -}
            where liftElem e = let !n3e = n3 e in (n3e, n3e, n3e)
        (q, 1) -> flatTop $ go (q - 1) t0 liftElem {- 2 3 2 -}
            where liftElem e = let !n2e = n2 e in (n2e, n3 e, n2e)
        (q, _) -> flatTop $ go (q - 1) t0 liftElem {- 3 3 2 -}
            where liftElem e = let !n3e = n3 e in (n3e, n3e, n2 e)
  where
    resSize = n0 * (size ft0 + 2)

    n2 e = Node2 2 e e
    n3 e = Node3 3 e e e

    go :: Int -> RepTop (Elem a) -> (Elem a -> (Node c, Node c, Node c))
       -> {- Int -> Int -> Int -> -} RepTop (Node c)
    go !n (RepTop l ft r) lift {- !sl !sm !sr -} =
        RepTop l1 (repEachMiddle n (RepMid l2 l3 ft r1 r2) lift) r3 {- sl sm sr -}
      where
        !(l1, l2, l3) = lift l
        !(r1, r2, r3) = lift r
    {-# INLINE go #-}

    -- invariants:
    -- - r1 has size sl; l2 and r2 have size sm; r3 has size sr
    -- - `lift` produces nodes of sizes sl, sm, sr with sl + n * sm + sr = n0
    -- - l2 and l3 are the left-most element replicated sm and sr times
    -- - r1 and r2 are the right-most element replicated sl and sm times
    repEachMiddle :: Int -> RepMid (Elem a) c
        -> (Elem a -> (Node c, Node c, Node c))
        {- -> Int -> Int -> Int -} -> FingerTree (Node c)
    repEachMiddle !n (RepMid l2 l3 ft r1 r2) lift {- !sl !sm !sr -} = case n of
        0 -> deep' (One l3)
                   (mapMulFT n0 lift0 ft)
                   (One r1)
        1 -> deep' (Two l2 l3)
                   (mapMulFT n0 lift1 ft)
                   (Two r1 r2)
        _ -> case n `quotRem` 3 of
            (q, 0) -> finR0 $ repEachMiddle (q - 1) midR0 liftR0 {- sl0 sm0 sr0 -}
            (q, 1) -> finR1 $ repEachMiddle (q - 1) midR1 liftR1 {- sl1 sm1 sr1 -}
            (q, _) -> finR2 $ repEachMiddle q       midR2 liftR2 {- sl2 sm2 sr2 -}
      where
        -- we could also pass these as arguments, see commented code above
        sl = size r1
        sm = size l2 -- = size r2
        sr = size l3

        -- the outermost constructor always has this size
        deep' = Deep (resSize - sr - sl)

        -- base cases: n = 0 or n = 1
        lift0 e = let !(l, _, r) = lift e in Node2 n0 l r
        lift1 e = let !(l, m, r) = lift e in Node3 n0 l m r

        -- n `mod` 3 = 0
        (!sl0, !sm0, !sr0) = (sl + sm, 3*sm, 2*sm + sr)
        liftR0 e = let !(l, m, r) = lift e in
                   (Node2 sl0 l m, Node3 sm0 m m m, Node3 sr0 m m r)
        midR0 = RepMid (Node3 sm0 l2 l2 l2)
                       (Node3 sr0 l2 l2 l3)
                       ft
                       (Node2 sl0 r1 r2)
                       (Node3 sm0 r2 r2 r2)
        finR0 ft' = deep' (One l2) ft' (Two r2 r2)

        -- n `mod` 3 = 1
        (!sl1, !sm1, !sr1) = (sl + 2*sm, 3*sm, 2*sm + sr)
        liftR1 e = let !(l, m, r) = lift e in
                   (Node3 sl1 l m m, Node3 sm1 m m m, Node3 sr1 m m r)
        midR1 = RepMid (Node3 sm1 l2 l2 l2)
                       (Node3 sr1 l2 l2 l3)
                       ft
                       (Node3 sl1 r1 r2 r2)
                       (Node3 sm1 r2 r2 r2)
        finR1 ft' = deep' (Two l2 l2) ft' (Two r2 r2)

        -- n `mod` 3 = 2
        (!sl2, !sm2, !sr2) = (sl + sm, 3*sm, sm + sr)
        liftR2 e = let !(l, m, r) = lift e in
                   (Node2 sl2 l m, Node3 sm2 m m m, Node2 sr2 m r)
        midR2 = RepMid (Node3 sm2 l2 l2 l2)
                       (Node2 sr2 l2 l3)
                       ft
                       (Node2 sl2 r1 r2)
                       (Node3 sm2 r2 r2 r2)
        finR2 ft' = deep' (One l2) ft' (One r2)

    -- flatten top part
    flatTop :: RepTop (Node a) -> FingerTree a
    flatTop (RepTop l ft r) = Deep resSize (nodeToDigit l) ft (nodeToDigit r)
    {-# INLINE flatTop #-}

    -- mapMul variant for top part
    mapMulTop :: Int -> (a -> b) -> RepTop a -> RepTop b
    mapMulTop n f (RepTop l ft r) = RepTop (f l) (mapMulFT n f ft) (f r)
    {-# INLINE mapMulTop #-}

{-# INLINE repEachTop #-}

I find it hard to stop. So I have floated out the common expression for the Deep constructors in my code, and eliminated the redundant sl, sm and sr parameters of repEachMiddle (Oh. That's not necessarily a win though, because recomputing those sizes involves a pattern match.).

The Wide trees are gone, because they were only used at the top level, so there's RepTop now instead.

Perhaps that RepTop abstraction is worth stealing? Note that your before2FT and before3FT compress down to flatTop (mapMulTop n0 (\e -> Node2 2 e e) (RepTop l ft r)) and flatTop (mapMulTop n0 (\e -> Node3 3 e e) (RepTop l ft r)), respectively, where l <| ft |> r is the initial finger tree. The resulting code should be pretty much the same after inlining.

Where would I use RepTop? I'd be happy to use whatever will make things less unclear!