[Haskell-cafe] Improving *> and >> for Data.Sequence

[Ross Paterson]

On Fri, Nov 21, 2014 at 02:00:16PM -0500, David Feuer wrote:
> To be precise, I *think* using the fromList approach for <*> makes us create O
> (n) thunks in order to extract the last element of the result. If we build the
> result inward, I *think* we can avoid this, getting the last element of the
> result in O(1) time and space. But my understanding of this data structure
> remains primitive.

This modification of the previous should do that.

mult :: Seq (a -> b) -> Seq a -> Seq b
mult sfs sxs = fromTwoLists (length sfs * length sxs) ys rev_ys
  where
    fs = toList sfs
    rev_fs = toRevList sfs
    xs = toList sxs
    rev_xs = toRevList sxs
    ys = [f x | f <- fs, x <- xs]
    rev_ys = [f x | f <- rev_fs, x <- rev_xs]

-- toRevList xs = toList (reverse xs)
toRevList :: Seq a -> [a]
toRevList = foldl (flip (:)) []

-- Build a tree lazy in the middle, from a list and its reverse.
--
-- fromTwoLists (length xs) xs (reverse xs) = fromList xs
--
-- Getting the kth element from either end involves forcing the lists
-- to length k.
fromTwoLists :: Int -> [a] -> [a] -> Seq a
fromTwoLists len_xs xs rev_xs =
    Seq $ mkTree2 len_xs 1 (map Elem xs) (map Elem rev_xs)

-- Construct a fingertree from the first n elements of xs.
-- The arguments must satisfy n <= length xs && rev_xs = reverse xs.
-- Each element of xs has the same size, provided as an argument.
mkTree2 :: Int -> Int -> [a] -> [a] -> FingerTree a
mkTree2 n size xs rev_xs
  | n == 0 = Empty
  | n == 1 = let [x1] = xs in Single x1
  | n <  6 = let
            nl = n `div` 2
            l = Data.List.take nl xs
            r = Data.List.take (n - nl) rev_xs
        in Deep totalSize (mkDigit l) Empty (mkRevDigit r)
  | otherwise = let
            size' = 3*size
            n' = (n-4) `div` 3
            digits = n - n'*3
            nl = digits `div` 2
            (l, xs') = Data.List.splitAt nl xs
            (r, rev_xs') = Data.List.splitAt (digits - nl) rev_xs
            nodes = mkNodes size' xs'
            rev_nodes = mkRevNodes size' rev_xs'
        in Deep totalSize (mkDigit l) (mkTree2 n' size' nodes rev_nodes) (mkRevDigit r)
  where
    totalSize = n*size

mkDigit :: [a] -> Digit a
mkDigit [x1] = One x1
mkDigit [x1, x2] = Two x1 x2
mkDigit [x1, x2, x3] = Three x1 x2 x3
mkDigit [x1, x2, x3, x4] = Four x1 x2 x3 x4

-- length xs <= 4 => mkRevDigit xs = mkDigit (reverse xs)
mkRevDigit :: [a] -> Digit a
mkRevDigit [x1] = One x1
mkRevDigit [x2, x1] = Two x1 x2
mkRevDigit [x3, x2, x1] = Three x1 x2 x3
mkRevDigit [x4, x3, x2, x1] = Four x1 x2 x3 x4

mkNodes :: Int -> [a] -> [Node a]
mkNodes size (x1:x2:x3:xs) = Node3 size x1 x2 x3:mkNodes size xs

-- length xs `mod` 3 == 0 =>
--    mkRevNodes size xs = reverse (mkNodes size (reverse xs))
mkRevNodes :: Int -> [a] -> [Node a]
mkRevNodes size (x3:x2:x1:xs) = Node3 size x1 x2 x3:mkRevNodes size xs