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 #-}

Could you explain the purpose of rigidifying the first/left argument here? We need to work with a rigid conceptualization of the second argument because we need to produce a replication function that makes a 2-3 tree, but why use a rigid first argument?

The rigidify is there because I implemented an intuition based on the 2-3 tree monad, and rigidify is how one gets an actual 2-3 tree.

This actually helped with the intial writing of the code--for example, I did have a flatRigid working on Rigid a.

In retrospect, rigidify does more work than needed here... the only property of rigidify that is used is that it produces digits of size 2 or 3 in the outermost constructor, and that could be relaxed to 2-4 at the expense of extra cases in flatRigidFT. This effect can be achieved the way you did: by splitting off the first and last elements from the finger tree.

I could change that... maybe I'll actually give it a go to see how closely the result would resemble your code.

@treeowl: Okay, I rewrote the code without rigidify and it's actually a bit simpler now, using a concept of a wide finger tree.

As far as I can see, the final <| and |> operations never hit the recursive case, but it's hard and probably worthless to express this fact in the code.

From here, one could start with adding the optimization that the left and the right part of the sequence can be partially produced on the fly. But I wonder how important that really is... does it matter whether the left-most element is available in O(1) or O(log(n)) time, where n is the amount of replication?

I don't think it's necessarily important to get that, but that's how we handle liftA2/<*> and *>, and follows the general approach of fmap, zipWith, and other bulk operations. Since we can get it, I think we'd need a good reason not to. As far as I know, the only remaining operation in the entire module that's not incrementally optimal is >< (append), and that was done intentionally for performance reasons.

Well it does add 50 lines of code and quite a bit of complexity.

Anyway... I've now implemented that and the code is starting to resemble yours a lot, especially once it gets to repEachMiddle.

Okay, some more refinement and testing done... I'll leave it at this. I don't see a reason to change the code you have. (I do like my own code better, but that's mostly because I wrote it; we have different styles.)

I do very much appreciate your attention to this. I think I'm going to stick with my own approach for now, unless there's a compelling reason to use this one. For the moment, at least, I understand mine at least somewhat better than yours.

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!