{-# LANGUAGE RankNTypes #-}

module Data.ByteString.Short.Encode (encodeUtf8, encodeUtf16LE) where

import Data.Bits
    ( shiftR, (.&.) )
import Data.ByteString.Short
    ( ShortByteString )
import Data.Char
    ( ord )
import Data.Word
    ( Word8 )

import qualified Data.ByteString.Short as BS
    ( pack )


encodeUtf8 :: String -> ShortByteString
encodeUtf8 :: String -> ShortByteString
encodeUtf8 = [Word8] -> ShortByteString
BS.pack ([Word8] -> ShortByteString)
-> (String -> [Word8]) -> String -> ShortByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> [Word8]
encode
  where
    encode :: String -> [Word8]
    encode :: String -> [Word8]
encode = (Char -> [Word8]) -> String -> [Word8]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Char -> [Word8]
encodeChar

    encodeChar :: Char -> [Word8]
    encodeChar :: Char -> [Word8]
encodeChar = (Int -> Word8) -> [Int] -> [Word8]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int] -> [Word8]) -> (Char -> [Int]) -> Char -> [Word8]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Int]
forall a. (Ord a, Num a, Bits a) => a -> [a]
go (Int -> [Int]) -> (Char -> Int) -> Char -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Int
ord
     where
      go :: a -> [a]
go a
oc
       | a
oc a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0x7f       = [a
oc]

       | a
oc a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0x7ff      = [ a
0xc0 a -> a -> a
forall a. Num a => a -> a -> a
+ (a
oc a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
6)
                            , a
0x80 a -> a -> a
forall a. Num a => a -> a -> a
+ a
oc a -> a -> a
forall a. Bits a => a -> a -> a
.&. a
0x3f
                            ]

       | a
oc a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
0xffff     = [ a
0xe0 a -> a -> a
forall a. Num a => a -> a -> a
+ (a
oc a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
12)
                            , a
0x80 a -> a -> a
forall a. Num a => a -> a -> a
+ ((a
oc a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
6) a -> a -> a
forall a. Bits a => a -> a -> a
.&. a
0x3f)
                            , a
0x80 a -> a -> a
forall a. Num a => a -> a -> a
+ a
oc a -> a -> a
forall a. Bits a => a -> a -> a
.&. a
0x3f
                            ]
       | Bool
otherwise        = [ a
0xf0 a -> a -> a
forall a. Num a => a -> a -> a
+ (a
oc a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
18)
                            , a
0x80 a -> a -> a
forall a. Num a => a -> a -> a
+ ((a
oc a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
12) a -> a -> a
forall a. Bits a => a -> a -> a
.&. a
0x3f)
                            , a
0x80 a -> a -> a
forall a. Num a => a -> a -> a
+ ((a
oc a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
6) a -> a -> a
forall a. Bits a => a -> a -> a
.&. a
0x3f)
                            , a
0x80 a -> a -> a
forall a. Num a => a -> a -> a
+ a
oc a -> a -> a
forall a. Bits a => a -> a -> a
.&. a
0x3f
                            ]
{-# INLINE encodeUtf8 #-}


encodeUtf16LE :: String -> ShortByteString
encodeUtf16LE :: String -> ShortByteString
encodeUtf16LE = [Word8] -> ShortByteString
BS.pack ([Word8] -> ShortByteString)
-> (String -> [Word8]) -> String -> ShortByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> [Word8]
encode
  where
    encode :: String -> [Word8]
    encode :: String -> [Word8]
encode = (Char -> [Word8]) -> String -> [Word8]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Char -> [Word8]
encodeChar

    encodeChar :: Char -> [Word8]
    encodeChar :: Char -> [Word8]
encodeChar = (Int -> Word8) -> [Int] -> [Word8]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int] -> [Word8]) -> (Char -> [Int]) -> Char -> [Word8]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Int]
forall a. (Ord a, Bits a, Num a) => a -> [a]
go (Int -> [Int]) -> (Char -> Int) -> Char -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Int
ord
      where
        go :: a -> [a]
go a
oc
          | a
oc a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0x10000 = [ a
oc, a
oc a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
8 ]
          | Bool
otherwise =
            let m :: a
m = a
oc a -> a -> a
forall a. Num a => a -> a -> a
- a
0x10000
            in [ a
m a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
10
               , (a
m a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
18) a -> a -> a
forall a. Num a => a -> a -> a
+ a
0xD8
               , a
m a -> a -> a
forall a. Bits a => a -> a -> a
.&. a
0x3FF 
               , ((a
m a -> a -> a
forall a. Bits a => a -> a -> a
.&. a
0x3FF) a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
8) a -> a -> a
forall a. Num a => a -> a -> a
+ a
0xDC ]
{-# INLINE encodeUtf16LE #-}