Brotli compression algorithm, translated from Python into Haskell

Question

I have converted some code in Python into its Haskell equivalent.

Python implementation:

def final_codes( bit_lengths, next_code ):
    tree_code = []
    for i in range(len(bit_lengths)):
        tree_code.append( next_code[ bit_lengths[i] ] )
        next_code[bit_lengths[i]] += 1
    return tree_code

#Inputs: 
bit_lengths = [ 3, 3, 3, 3, 3, 2, 4, 4 ]
next_code = [0,0,0,2,14]
print final_codes( bit_lengths, next_code )
[2, 3, 4, 5, 6, 0, 14, 15]

Haskell implementation:

replaceNth :: (Eq a, Num a) => a -> a1 -> [a1] -> [a1]
replaceNth n newVal (x:xs)       # Picked this function from somewhere on stackoverflow.
    | n == 0 = newVal:xs
    | otherwise = x:replaceNth (n-1) newVal xs

final_codes' :: (Eq a, Num a, Num a1) => [a1] -> [Int] -> [a1] -> a -> [a1]
final_codes' lst bit_lngths next_cd 0 = lst
final_codes' lst bit_lngths next_cd n = final_codes' new_lst bit_lngths new_next_code (n-1)
    where tlen = bit_lngths !! (length lst)
          new_lst = (next_cd !! tlen) : lst
          new_next_code = replaceNth tlen ((next_cd !! tlen) + 1) next_cd  -- INEFFICIENT!

final_codes :: Num a => [Int] -> [a] -> [a]          
final_codes bit_lengths next_code = reverse $ final_codes' [] bit_lengths next_code (length bit_lengths)

In the last where statement in final_codes' I create a new_next_code list by incrementing the value at position tlen by 1. This makes the algorithm \$O(N^2)\$ in comparison to the \$O(N)\$ implementation in Python.

This approach works but it feels a bit contrived. Any suggestions on how I could improve this code and its performance?

Note: I'd like to avoid using the lens library for now (I'm still a novice). I found this Stack Overflow link on replacing a single element in a list; but most approaches are \$O(N)\$. I haven't yet dug into the versions that use lens.

Answer 1

don't use lists

Almost any non-trivial operation on lists is O(n) or worse. Python is using mutable arrays. Haskell also mutable arrays, so why not use mutable arrays?

The only problem with Haskell's arrays is that there are so many choices to make:

• pure vs. mutable
• boxed vs. unboxed
• array vs. vector
• run in the ST monad or IO?

Here is the routine coded in Haskell which uses both pure and mutable unboxed vectors. It could be improved, but this is basically how you would go about it.

(code updated)

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 )
-}

Answer 2

Simplest solution would be to use IntMap instead of mutable vector (with logarithmic overhead):

import Data.IntMap (IntMap)
import qualified Data.IntMap as Map

finalCodes :: [Int] -> IntMap Int -> [Int]
finalCodes [] _ = []
finalCodes (bitLen:bitLengths) nextCodes
  = let (Just x, nextCodes')
          = Map.updateLookupWithKey (const $ Just . (+1)) bitLen nextCodes
    in  x : finalCodes bitLengths nextCodes'


bitLengths = [3, 3, 3, 3, 3, 2, 4, 4]
nextCodes = [0, 0, 0, 2, 14]
test = finalCodes bitLengths $ Map.fromList $ zip [0..] nextCodes

A bit of monad hackery is required to use Data.Vector.Unboxed.Mutable instead:

import Control.Monad (forM)
import qualified Data.Vector.Unboxed as V
import qualified Data.Vector.Unboxed.Mutable as MV

finalCodes :: [Int] -> [Int] -> IO [Int]
finalCodes bitLengths nextCodesLst = do
  nextCodes <- V.thaw $ V.fromList nextCodesLst
  forM bitLengths $ \bitLen -> do
    x <- MV.read nextCodes bitLen
    MV.write nextCodes bitLen (x+1)
    return x

Btw, please note, it is more common to see camel case in Haskell code: https://wiki.haskell.org/Camel_case