Skip to content

Instantly share code, notes, and snippets.

@int-e
Last active April 14, 2020 17:28
Embed
What would you like to do?
replicateEach
-- `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 #-}
-- `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 #-}
@int-e
Copy link
Author

int-e commented Apr 13, 2020

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.

@int-e
Copy link
Author

int-e commented Apr 13, 2020

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.)

@treeowl
Copy link

treeowl commented Apr 14, 2020

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.

@int-e
Copy link
Author

int-e commented Apr 14, 2020

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.

@treeowl
Copy link

treeowl commented Apr 14, 2020

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment