Computing fixed-width monoidal sliding windows with chunked partial sums

Description.md

In Brent Yorgey's post about monoidal sliding windows he used a queue to represent the fixed-size sliding window.

Due to the fixed size, however, the append stack was reversed to become a new pop stack for the queue at very regular intervals. Intervals of exactly the width of the fixed window.

In effect then, what the queue was computing was prefix and suffix monoidal sums, not over the entire list, but over discrete chunks of the list the one smaller than the size of the fixed window.

For example, say the window size is 5, and our input list is [a0, a1, a2, ..., a21].

If we compute the list of prefix sums of chunks of size 4:

 chunkedPrefixSums =
   [ a0, a0 <> a1, a0 <> a1 <> a2, a0 <> a1 <> a2 <> a3
   , a4, a4 <> a5, a4 <> a5 <> a6, a4 <> a5 <> a6 <> a7
   , a8, a8 <> a9, a8 <> a9 <> a10, a8 <> a9 <> a10 <> a11
   , ...
   , a16, a16 <> a17, a16 <> a17 <> a18, a16 <> a17 <> a18 <> a19
   , a20, a20 <> a21
   ]

and the list of suffix sums of chunks of size 4

 chunkedSuffixSums = 
   [ a0 <> a1 <> a2 <> a3, a1 <> a2 <> a3, a2 <> a3, a3
   , a4 <> a5 <> a6 <> a7, a5 <> a6 <> a7, a6 <> a7, a7
   , ...
   , a12 <> a13 <> a14 <> a15, a13 <> a14 <> a15, a14 <> a15, a15
   , a16 <> a17 <> a18 <> a19, a17 <> a18 <> a19, a18 <> a19, a19
   , a20 <> a21, a21
   ]

Then we can use them to compute the monoidal sums of all windows of size 5 by zipping them together carefully

 zipWith (<>) chunkedSuffixSums (drop 4 chunkedPrefixSums) =
   [ (a0 <> a1 <> a2 <> a3) <> a4
   , (a1 <> a2 <> a3) <> (a4 <> a5)
   , (a2 <> a3) <> (a4 <> a5 <> a6)
   , ...
   , a15 <> (a16 <> a17 <> a18 <> a19)
   , (a16 <> a17 <> a18 <> a19) <> a20
   , (a17 <> a18 <> a19) <> (a20 <> a21)
   ]

FixedWindow.hs

module FixedWindow where

-- | Compute all the monoidal sums of a sliding window of the input with
-- exactly the given number of elements
--
--    >>> windowSums 4 ["a","b","c","d","e","f","g","h","i","j"]
--    ["abcd","bcde","cdef","defg","efgh","fghi","ghij"]
--
--    >>> take 30 $ map getSum $ windowSums 4 $ map Sum $ cycle [1,1,1,1,0,0,0,0]
--    [4,3,2,1,0,1,2,3,4,3,2,1,0,1,2,3,4,3,2,1,0,1,2,3,4,3,2,1,0,1]
windowSums :: Monoid m => Word -> [m] -> [m]
windowSums 0 _ = repeat mempty
windowSums 1 xs = xs
windowSums n xs = zipWith (<>)
  (chunkedSuffixSums (n-1) xs)
  (drop (fromEnum n-1) $ chunkedPrefixSums (n-1) xs)

-- | Compute the suffix sums in each chunk of the given length in the input
--
--    >>> chunkedSuffixSums 3 ["a","b","c","d","e","f","g","h","i","j"]
--    ["abc","bc","c","def","ef","f","ghi","hi","i","j"]
--
--    >>> take 30 $ map getSum $ chunkedSuffixSums 3 $ map Sum $ cycle [1,1,1,1,0,0,0,0]
--    [3,2,1,1,0,0,1,1,1,3,2,1,0,0,0,2,2,1,2,1,0,0,0,0,3,2,1,1,0,0]
chunkedSuffixSums :: Monoid m => Word -> [m] -> [m]
chunkedSuffixSums 0 = \_ -> repeat mempty
chunkedSuffixSums maxSize = safeInit . scanr step mempty . zip chunkEnds where
  chunkEnds :: [Bool]
  chunkEnds = cycle $ replicate (fromEnum maxSize - 1) False <> [True]

  step :: Semigroup m => (Bool, m) -> m -> m
  step (chunkEnd, m) total
    | chunkEnd = m
    | otherwise = m <> total

  safeInit :: [a] -> [a]
  safeInit [] = []
  safeInit (a:as) = loop a as where
    loop _ [] = []
    loop b (a:as) = b : loop a as

-- | Compute the prefix sums in each chunk of the given length in the input
--
--    >>> chunkedPrefixSums 3 ["a","b","c","d","e","f","g","h","i","j"]
--    ["a","ab","abc","d","de","def","g","gh","ghi","j"]
--
--    >>> take 30 $ map getSum $ chunkedPrefixSums 3 $ map Sum $ cycle [1,1,1,1,0,0,0,0]
--    [1,2,3,1,1,1,0,0,1,1,2,3,0,0,0,0,1,2,1,2,2,0,0,0,1,2,3,1,1,1]
chunkedPrefixSums :: Monoid m => Word -> [m] -> [m]
chunkedPrefixSums 0 = \_ -> repeat mempty
chunkedPrefixSums maxSize = safeTail . scanl step mempty . zip chunkStarts where
  chunkStarts :: [Bool]
  chunkStarts = cycle $ True : replicate (fromEnum maxSize - 1) False

  step :: Semigroup m => m -> (Bool, m) -> m
  step total (chunkStart, m) 
    | chunkStart = m
    | otherwise = total <> m

  safeTail :: [a] -> [a]
  safeTail [] = []
  safeTail (_: as) = as