-- |

-- Module      : Data.Memory.Encoding.Base32

-- License     : BSD-style

-- Maintainer  : Nicolas DI PRIMA <nicolas@di-prima.fr>

-- Stability   : experimental

-- Portability : unknown

--

-- Base32

--

{-# LANGUAGE MagicHash #-}

{-# LANGUAGE UnboxedTuples #-}

{-# LANGUAGE OverloadedStrings #-}

{-# LANGUAGE BangPatterns #-}

{-# LANGUAGE Rank2Types #-}

module Data.Memory.Encoding.Base32

    ( toBase32

    , unBase32Length

    , fromBase32

    ) where

import           Data.Memory.Internal.Compat

import           Data.Memory.Internal.CompatPrim

import           Data.Word

import           Data.Bits ((.|.))

import           GHC.Prim

import           GHC.Word

import           Control.Monad

import           Foreign.Storable

import           Foreign.Ptr (Ptr)

-- | Transform a number of bytes pointed by.@src in the base32 binary representation in @dst

--

-- destination memory need to be of correct size, otherwise it will lead

-- to really bad things.

toBase32 :: Ptr Word8 -- ^ input

         -> Ptr Word8 -- ^ output

         -> Int       -- ^ input len

         -> IO ()

toBase32 dst src len = loop 0 0

  where

    eqChar :: Word8

    eqChar = 0x3d

    peekOrZero :: Int -> IO Word8

    peekOrZero i

        | i >= len  = return 0

        | otherwise = peekByteOff src i

    pokeOrPadding :: Int -- for the test

                  -> Int -- src index

                  -> Word8 -- the value

                  -> IO ()

    pokeOrPadding i di v

        | i <  len  = pokeByteOff dst di v

        | otherwise = pokeByteOff dst di eqChar

    loop :: Int -- index input

         -> Int -- index output

         -> IO ()

    loop i di

        | i >  len  = return ()

        | otherwise = do

            i1 <- peekByteOff src i

            i2 <- peekOrZero (i + 1)

            i3 <- peekOrZero (i + 2)

            i4 <- peekOrZero (i + 3)

            i5 <- peekOrZero (i + 4)

            let (o1,o2,o3,o4,o5,o6,o7,o8) = toBase32Per5Bytes (i1, i2, i3, i4, i5)

            pokeByteOff dst di o1

            pokeByteOff dst (di + 1) o2

            pokeOrPadding (i + 1) (di + 2) o3

            pokeOrPadding (i + 1) (di + 3) o4

            pokeOrPadding (i + 2) (di + 4) o5

            pokeOrPadding (i + 3) (di + 5) o6

            pokeOrPadding (i + 3) (di + 6) o7

            pokeOrPadding (i + 4) (di + 7) o8

            loop (i+5) (di+8)

toBase32Per5Bytes :: (Word8, Word8, Word8, Word8, Word8)

                  -> (Word8, Word8, Word8, Word8, Word8, Word8, Word8, Word8)

toBase32Per5Bytes (W8# i1, W8# i2, W8# i3, W8# i4, W8# i5) =

    (index o1, index o2, index o3, index o4, index o5, index o6, index o7, index o8)

  where

    -- 1111 1000 >> 3

    !o1 =     (uncheckedShiftRL# (and# i1 0xF8##) 3#)

    -- 0000 0111 << 2 | 1100 0000 >> 6

    !o2 = or# (uncheckedShiftL#  (and# i1 0x07##) 2#) (uncheckedShiftRL# (and# i2 0xC0##) 6#)

    -- 0011 1110 >> 1

    !o3 =     (uncheckedShiftRL# (and# i2 0x3E##) 1#)

    -- 0000 0001 << 4 | 1111 0000 >> 4

    !o4 = or# (uncheckedShiftL#  (and# i2 0x01##) 4#) (uncheckedShiftRL# (and# i3 0xF0##) 4#)

    -- 0000 1111 << 1 | 1000 0000 >> 7

    !o5 = or# (uncheckedShiftL#  (and# i3 0x0F##) 1#) (uncheckedShiftRL# (and# i4 0x80##) 7#)

    -- 0111 1100 >> 2

    !o6 =     (uncheckedShiftRL# (and# i4 0x7C##) 2#)

    -- 0000 0011 << 3 | 1110 0000 >> 5

    !o7 = or# (uncheckedShiftL#  (and# i4 0x03##) 3#) (uncheckedShiftRL# (and# i5 0xE0##) 5#)

    -- 0001 1111

    !o8 =     ((and# i5 0x1F##))

    !set = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567"#

    index :: Word# -> Word8

    index idx = W8# (indexWord8OffAddr# set (word2Int# idx))

-- | Get the length needed for the destination buffer for a base32 decoding.

--

-- if the length is not a multiple of 8, Nothing is returned

unBase32Length :: Ptr Word8 -> Int -> IO (Maybe Int)

unBase32Length src len

    | (len `mod` 8) /= 0 = return Nothing

    | otherwise          = do

        last1Byte <- peekByteOff src (len - 1)

        last2Byte <- peekByteOff src (len - 2)

        last3Byte <- peekByteOff src (len - 3)

        last4Byte <- peekByteOff src (len - 4)

        last5Byte <- peekByteOff src (len - 5)

        last6Byte <- peekByteOff src (len - 6)

        let dstLen = caseByte last1Byte last2Byte last3Byte last4Byte last5Byte last6Byte

        return $ Just $ (len `div` 8) * 5 - dstLen

  where

    caseByte :: Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> Int

    caseByte last1 last2 last3 last4 last5 last6

        | last6 == eqAscii = 4

        | last5 == eqAscii = 3 -- error this padding is not expected (error will be detected in fromBase32)

        | last4 == eqAscii = 3

        | last3 == eqAscii = 2

        | last2 == eqAscii = 1 -- error this padding is not expected (error will be detected in fromBase32)

        | last1 == eqAscii = 1

        | otherwise        = 0

    eqAscii :: Word8

    eqAscii = 0x3D

-- | convert from base32 in @src to binary in @dst, using the number of bytes specified

--

-- the user should use unBase32Length to compute the correct length, or check that

-- the length specification is proper. no check is done here.

fromBase32 :: Ptr Word8 -> Ptr Word8 -> Int -> IO (Maybe Int)

fromBase32 dst src len

    | len == 0  = return Nothing

    | otherwise = loop 0 0

  where

    loop :: Int -- the index dst

         -> Int -- the index src

         -> IO (Maybe Int)

    loop di i

        | i == (len - 8) = do

            i1 <- peekByteOff src i

            i2 <- peekByteOff src (i + 1)

            i3 <- peekByteOff src (i + 2)

            i4 <- peekByteOff src (i + 3)

            i5 <- peekByteOff src (i + 4)

            i6 <- peekByteOff src (i + 5)

            i7 <- peekByteOff src (i + 6)

            i8 <- peekByteOff src (i + 7)

            let (nbBytes, i3', i4', i5', i6', i7', i8') =

                    case (i3, i4, i5, i6, i7, i8) of

                        (0x3D, 0x3D, 0x3D, 0x3D, 0x3D, 0x3D) -> (6, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41)

                        (0x3D, _   , _   , _   , _   , _   ) -> (0, i3, i4, i5, i6, i7, i8) -- invalid

                        (_   , 0x3D, 0x3D, 0x3D, 0x3D, 0x3D) -> (5, i3  , 0x41, 0x41, 0x41, 0x41, 0x41)

                        (_   , 0x3D, _   , _   , _   , _   ) -> (0, i3, i4, i5, i6, i7, i8) -- invalid

                        (_   , _   , 0x3D, 0x3D, 0x3D, 0x3D) -> (4, i3  , i4  , 0x41, 0x41, 0x41, 0x41)

                        (_   , _   , 0x3D, _   , _   , _   ) -> (0, i3, i4, i5, i6, i7, i8) -- invalid

                        (_   , _   , _   , 0x3D, 0x3D, 0x3D) -> (3, i3  , i4  , i5  , 0x41, 0x41, 0x41)

                        (_   , _   , _   , 0x3D, _   , _   ) -> (0, i3, i4, i5, i6, i7, i8) -- invalid

                        (_   , _   , _   , _   , 0x3D, 0x3D) -> (2, i3  , i4  , i5  , i6  , 0x41, 0x41)

                        (_   , _   , _   , _   , 0x3D, _   ) -> (0, i3, i4, i5, i6, i7, i8) -- invalid

                        (_   , _   , _   , _   , _   , 0x3D) -> (1, i3  , i4  , i5  , i6  , i7  , 0x41)

                        (_   , _   , _   , _   , _   , _   ) -> (0 :: Int, i3, i4, i5, i6, i7, i8)

            case fromBase32Per8Bytes (i1, i2, i3', i4', i5', i6', i7', i8') of

                Left  ofs                  -> return $ Just (i + ofs)

                Right (o1, o2, o3, o4, o5) -> do

                    pokeByteOff dst  di    o1

                    pokeByteOff dst (di+1) o2

                    when (nbBytes < 5) $ pokeByteOff dst (di+2) o3

                    when (nbBytes < 4) $ pokeByteOff dst (di+3) o4

                    when (nbBytes < 2) $ pokeByteOff dst (di+4) o5

                    return Nothing

        | otherwise = do

            i1 <- peekByteOff src i

            i2 <- peekByteOff src (i + 1)

            i3 <- peekByteOff src (i + 2)

            i4 <- peekByteOff src (i + 3)

            i5 <- peekByteOff src (i + 4)

            i6 <- peekByteOff src (i + 5)

            i7 <- peekByteOff src (i + 6)

            i8 <- peekByteOff src (i + 7)

            case fromBase32Per8Bytes (i1, i2, i3, i4, i5, i6, i7, i8) of

                Left  ofs                  -> return $ Just (i + ofs)

                Right (o1, o2, o3, o4, o5) -> do

                    pokeByteOff dst  di    o1

                    pokeByteOff dst (di+1) o2

                    pokeByteOff dst (di+2) o3

                    pokeByteOff dst (di+3) o4

                    pokeByteOff dst (di+4) o5

                    loop (di+5) (i+8)

fromBase32Per8Bytes :: (Word8, Word8, Word8, Word8, Word8, Word8, Word8, Word8)

                    -> Either Int (Word8, Word8, Word8, Word8, Word8)

fromBase32Per8Bytes (i1, i2, i3, i4, i5, i6, i7, i8) =

    case (rset i1, rset i2, rset i3, rset i4, rset i5, rset i6, rset i7, rset i8) of

        (0xFF, _   , _   , _   , _   , _   , _   , _   ) -> Left 0

        (_   , 0xFF, _   , _   , _   , _   , _   , _   ) -> Left 1

        (_   , _   , 0xFF, _   , _   , _   , _   , _   ) -> Left 2

        (_   , _   , _   , 0xFF, _   , _   , _   , _   ) -> Left 3

        (_   , _   , _   , _   , 0xFF, _   , _   , _   ) -> Left 4

        (_   , _   , _   , _   , _   , 0xFF, _   , _   ) -> Left 5

        (_   , _   , _   , _   , _   , _   , 0xFF, _   ) -> Left 6

        (_   , _   , _   , _   , _   , _   , _   , 0xFF) -> Left 7

        (ri1 , ri2 , ri3 , ri4 , ri5 , ri6 , ri7 , ri8 ) ->

                -- 0001 1111 << 3 | 0001 11xx >> 2

            let o1 = (ri1 `unsafeShiftL` 3) .|. (ri2 `unsafeShiftR` 2)

                -- 000x xx11 << 6 | 0001 1111 << 1 | 0001 xxxx >> 4

                o2 = (ri2 `unsafeShiftL` 6) .|. (ri3 `unsafeShiftL` 1) .|. (ri4 `unsafeShiftR` 4)

                -- 000x 1111 << 4 | 0001 111x >> 1

                o3 = (ri4 `unsafeShiftL` 4) .|. (ri5 `unsafeShiftR` 1)

                -- 000x xxx1 << 7 | 0001 1111 << 2 | 0001 1xxx >> 3

                o4 = (ri5 `unsafeShiftL` 7) .|. (ri6 `unsafeShiftL` 2) .|. (ri7 `unsafeShiftR` 3)

                -- 000x x111 << 5 | 0001 1111

                o5 = (ri7 `unsafeShiftL` 5) .|. ri8

             in Right (o1, o2, o3, o4, o5)

  where

    rset :: Word8 -> Word8

    rset (W8# w)

        | booleanPrim (w `leWord#` 0xff##) = W8# (indexWord8OffAddr# rsetTable (word2Int# w))

        | otherwise                        = 0xff

    !rsetTable = "\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\

                 \\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\

                 \\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\

                 \\xFF\xFF\x1A\x1B\x1C\x1D\x1E\x1F\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\

                 \\xFF\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0A\x0B\x0C\x0D\x0E\

                 \\x0F\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\xFF\xFF\xFF\xFF\xFF\

                 \\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\

                 \\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\

                 \\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\

                 \\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\

                 \\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\

                 \\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\

                 \\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\

                 \\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\

                 \\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\

                 \\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF\xFF"#