Implement Insertion Sort in Haskell to understand ST mutation

insertion_sort.hs

module Main where

import Control.Monad.ST
import qualified Data.STRef as STR
import qualified Data.Array.ST as STArray
import qualified Data.Array as Arr
import qualified Data.Array.MArray as MA


untilM :: (Monad m) => m Bool -> m a -> m ()
untilM p f = go
  where go = do
          x <- p
          if x
            then return ()
            else f >> go

insertionSort unordered = ordered
  where
    ordered = STArray.runSTArray $ do
      indexRef <- STR.newSTRef 1 :: ST s (STR.STRef s Integer)
      
      a <- MA.thaw unordered
      untilM  (do
                  index <- STR.readSTRef indexRef
                  return (index == 5)
              ) (do
                    elements <- MA.getElems a
                    index <- STR.readSTRef indexRef
                    elemAtIndex <- MA.readArray a index
                    nextElement <- MA.readArray a $ index + 1
                    if (elemAtIndex < nextElement)
                      then STR.writeSTRef indexRef (index + 1)
                      else do MA.writeArray a index nextElement
                              MA.writeArray a (index + 1) elemAtIndex
                              STR.writeSTRef indexRef $ max 1 (index - 1)
                )
      return a

main = print $ insertionSort $ Arr.listArray (1, 5) [5,4,3,2,1]