b3.hs

b3.hs

{-
Response to http://lists.samba.org/archive/linux/2013-November/032604.html
and http://lists.samba.org/archive/linux/2013-November/032610.html

Author: Carlo Hamalainen <carlo@carlo-hamalainen.net>
-}

import Control.Monad
import Data.Binary.Get
import Data.Bits
import Data.Char
import Data.List.Split
import Data.Maybe
import Data.Word

import qualified Data.ByteString.Lazy as BSL

-- Extract the 2-bit pairs from a word.
pairs :: Word8 -> BSL.ByteString
pairs x = BSL.concat $ map BSL.singleton $ reverse [3 .&. shiftR x (2*i) | i <- [0..3]]

-- Corrector keeps the first bit of each bit pair
-- that is not 0b00 or 0b11.
correctPair :: Word8 -> BSL.ByteString
correctPair 1 = BSL.singleton 0
correctPair 2 = BSL.singleton 1
correctPair _ = BSL.empty

-- Nicked from http://stackoverflow.com/a/13219751
chunkBytes :: Int -> BSL.ByteString -> [BSL.ByteString]
chunkBytes n xs = chunk' i xs
  where
    chunk' n xs
        | BSL.null xs = []
        | otherwise = a : chunk' n b where (a,b) = BSL.splitAt n xs
    i = ceiling (fromIntegral (BSL.length xs) / fromIntegral n)

-- Attempt to assemble a byte from 8 bits.
-- assert length bits == 8
assembleByte bits = foldl (.|.) 0 [ shiftL (BSL.index bits (fromIntegral n)) (7 - n) | n <- [0..7]]

-- Assemble as many bytes as possible from a list of bits.
assembleBytes bits = f BSL.empty bits
    where f acc bits = if BSL.length bits < 8
                        then acc
                        else f (BSL.snoc acc $ assembleByte (BSL.take 8 bits)) (BSL.drop 8 bits)

-- Remove bias in a sequence of words.
extractor = BSL.pack . map assembleByte . chunkBytes 8 . BSL.concatMap correctPair . BSL.concatMap pairs

main = do
    c <- BSL.getContents
    BSL.putStr $ extractor c