Waterfall in ST

waterfall.hs

module Waterfall.Base where

import           Control.Monad.ST
import           Data.Foldable               (for_)
import           Data.STRef
import qualified Data.Vector.Unboxed         as U
import qualified Data.Vector.Unboxed.Mutable as MU

rainfallBase :: [Int] -> Int
rainfallBase xs = sum (zipWith (-) mins xs)
  where
    mins = zipWith min maxl maxr
    maxl = scanl1 max xs
    maxr = scanr1 max xs

rainfallUnboxedVector :: U.Vector Int -> Int
rainfallUnboxedVector xs = U.sum (U.zipWith (-) mins xs)
  where
    mins = U.zipWith min maxl maxr
    maxl = U.scanl1' max xs
    maxr = U.scanr1' max xs

rainfallST :: U.Vector Int -> Int
rainfallST xs = runST $ do
    let len = U.length xs
    towers <- MU.new len
    mx <- newSTRef 0
    sum' <- newSTRef 0

    for_ [0..len] $ \i -> do
        let t = U.unsafeIndex xs i
        modifySTRef' mx (max t)
        x <- readSTRef mx
        MU.unsafeWrite towers i x

    writeSTRef mx 0
    for_ [len-1, len-2 .. 0] $ \i -> do
        let t = U.unsafeIndex xs i
        modifySTRef' mx (max t)
        x <- readSTRef mx
        m <- MU.read towers i
        modifySTRef' sum' (\s -> s + (min x m - t))

    readSTRef sum'