part of the brutal compression algorithm in Haskell

brotli.hs

module Lib
where

import Control.Monad
import Control.Monad.Primitive
import Control.Monad.ST
import qualified Data.Vector.Unboxed.Mutable as UVM
import qualified Data.Vector.Unboxed as UV

calculate_codes :: PrimMonad m => UV.Vector Int -> UV.Vector Int -> m (UV.Vector Int)
calculate_codes tree codes = do
  next_codes <- UV.thaw codes
  UV.generateM (UV.length tree) $ \i -> do
    let x = tree UV.! i
    z <- if x /= 0
           then do y <- UVM.read next_codes x
                   UVM.write next_codes x (y+1)
                   return y
           else return 0
    return z

test = do
  let tree = UV.fromList [ 3, 3, 3, 3, 3, 2, 4, 4 ]
      codes = UV.fromList [ 0, 0, 0, 2, 14 ]
      next = runST $ calculate_codes tree codes
  print next

{-
tree      = [ 3, 3, 3, 3, 3, 2, 4, 4 ]
next_code = [0,0,0,2,14]

def calculate_codes( tree, codes ):
    tree_code = {}
    next_code = copy.copy( codes )
    for i in range(len(tree)):
        if tree[i] != 0:
            tree_code[i] = next_code[ tree[i] ]
            next_code[tree[i]] += 1             -- This "effectful" line is throwing me off.
    return tree_code

print calculate_codes( tree, next_code )
-}