module Country
( Country
, encodeNumeric
, decodeNumeric
, encodeEnglish
, decode
, alphaTwoUpper
, alphaThreeUpper
, alphaThreeLower
, alphaTwoLower
, decodeAlphaTwo
, decodeAlphaThree
) where
import Country.Unsafe (Country(..))
import Country.Unexposed.Encode.English (countryNameQuads)
import Country.Unexposed.ExtraNames (extraNames)
import Country.Unexposed.Names (englishCountryNamesText,numberOfPossibleCodes)
import Country.Unexposed.Enumerate (enumeratedCountries)
import Data.Text (Text)
import Data.ByteString (ByteString)
import Data.Word (Word16,Word8)
import Data.Primitive (indexArray,newArray,unsafeFreezeArray,writeArray,
writeByteArray,indexByteArray,unsafeFreezeByteArray,newByteArray)
import Data.HashMap.Strict (HashMap)
import Data.Primitive.Array (Array(..))
import Data.Primitive.ByteArray (ByteArray(..))
import GHC.Prim (sizeofByteArray#,sizeofArray#)
import GHC.Int (Int(..))
import Control.Monad.ST (runST)
import Control.Monad
import Data.Char (ord,chr,toLower)
import Data.Bits (unsafeShiftL,unsafeShiftR)
import qualified Data.List as L
import qualified Data.HashMap.Strict as HM
import qualified Data.Text as T
import qualified Data.Text.Array as TA
import qualified Data.Text.Internal as TI
encodeNumeric :: Country -> Word16
encodeNumeric (Country n) = n
decodeNumeric :: Word16 -> Maybe Country
decodeNumeric n = if n < 1000 && indexByteArray numericValidities (word16ToInt n) == (1 :: Word8)
then Just (Country n)
else Nothing
encodeEnglish :: Country -> Text
encodeEnglish (Country n) = indexArray englishCountryNamesText (word16ToInt n)
alphaTwoUpper :: Country -> Text
alphaTwoUpper c = TI.text allAlphaTwoUpper (timesTwo (indexOfCountry c)) 2
alphaThreeUpper :: Country -> Text
alphaThreeUpper c = TI.text allAlphaThreeUpper (timesThree (indexOfCountry c)) 3
alphaTwoLower :: Country -> Text
alphaTwoLower c = TI.text allAlphaTwoLower (timesTwo (indexOfCountry c)) 2
alphaThreeLower :: Country -> Text
alphaThreeLower c = TI.text allAlphaThreeLower (timesThree (indexOfCountry c)) 3
decodeAlphaTwo :: Text -> Maybe Country
decodeAlphaTwo = flip HM.lookup alphaTwoHashMap
decodeAlphaThree :: Text -> Maybe Country
decodeAlphaThree = flip HM.lookup alphaThreeHashMap
alphaTwoHashMap :: HashMap Text Country
alphaTwoHashMap = L.foldl'
(\hm (countryNum,_,(c1,c2),_) ->
HM.insert (T.pack [c1,c2]) (Country countryNum)
$ HM.insert (T.pack [toLower c1, toLower c2]) (Country countryNum)
$ hm
)
HM.empty countryNameQuads
alphaThreeHashMap :: HashMap Text Country
alphaThreeHashMap = L.foldl'
(\hm (countryNum,_,_,(c1,c2,c3)) ->
HM.insert (T.pack [c1,c2,c3]) (Country countryNum)
$ HM.insert (T.pack [toLower c1, toLower c2, toLower c3]) (Country countryNum)
$ hm
)
HM.empty countryNameQuads
half :: Int -> Int
half x = unsafeShiftR x 1
timesTwo :: Int -> Int
timesTwo x = unsafeShiftL x 1
timesThree :: Int -> Int
timesThree x = x * 3
decode :: Text -> Maybe Country
decode = flip HM.lookup decodeMap
word16ToInt :: Word16 -> Int
word16ToInt = fromIntegral
intToWord16 :: Int -> Word16
intToWord16 = fromIntegral
charToWord16 :: Char -> Word16
charToWord16 = fromIntegral . ord
word16ToChar :: Word16 -> Char
word16ToChar = chr . fromIntegral
decodeMap :: HashMap Text Country
decodeMap =
let baseMap = HM.union alphaTwoHashMap alphaThreeHashMap
hm1 = L.foldl' (\hm (country,name) -> HM.insert name country hm) baseMap extraNames
hm2 = L.foldl' (\hm (countryNum,name,_,_) -> HM.insert name (Country countryNum) hm) hm1 countryNameQuads
in hm2
arrayFoldl' :: (a -> b -> a) -> a -> Array b -> a
arrayFoldl' f z a = go 0 z
where
go i !acc | i < sizeofArray a = go (i+1) (f acc $ indexArray a i)
| otherwise = acc
sizeofArray :: Array a -> Int
sizeofArray (Array a) = I# (sizeofArray# a)
numberOfCountries :: Int
numberOfCountries = length countryNameQuads
numericValidities :: ByteArray
numericValidities = runST $ do
m <- newByteArray numberOfPossibleCodes
let clear !ix = if ix < numberOfPossibleCodes
then writeByteArray m ix (0 :: Word8)
else return ()
clear 0
forM_ countryNameQuads $ \(n,_,_,_) -> do
writeByteArray m (word16ToInt n) (1 :: Word8)
unsafeFreezeByteArray m
positions :: ByteArray
positions = runST $ do
m <- newByteArray (timesTwo numberOfPossibleCodes)
forM_ (zip (enumFrom (0 :: Word16)) countryNameQuads) $ \(ix,(n,_,_,_)) -> do
writeByteArray m (word16ToInt n) ix
unsafeFreezeByteArray m
indexOfCountry :: Country -> Int
indexOfCountry (Country n) =
word16ToInt (indexByteArray positions (word16ToInt n))
allAlphaTwoUpper :: TA.Array
allAlphaTwoUpper = TA.run $ do
m <- TA.new (timesTwo numberOfCountries)
forM_ countryNameQuads $ \(n,_,(a1,a2),_) -> do
let ix = timesTwo (indexOfCountry (Country n))
TA.unsafeWrite m ix (charToWord16 a1)
TA.unsafeWrite m (ix + 1) (charToWord16 a2)
return m
allAlphaThreeUpper :: TA.Array
allAlphaThreeUpper = TA.run $ do
m <- TA.new (timesThree numberOfCountries)
forM_ countryNameQuads $ \(n,_,_,(a1,a2,a3)) -> do
let ix = timesThree (indexOfCountry (Country n))
TA.unsafeWrite m ix (charToWord16 a1)
TA.unsafeWrite m (ix + 1) (charToWord16 a2)
TA.unsafeWrite m (ix + 2) (charToWord16 a3)
return m
allAlphaThreeLower :: TA.Array
allAlphaThreeLower = mapTextArray toLower allAlphaThreeUpper
allAlphaTwoLower :: TA.Array
allAlphaTwoLower = mapTextArray toLower allAlphaTwoUpper
mapTextArray :: (Char -> Char) -> TA.Array -> TA.Array
mapTextArray f a@(TA.Array inner) = TA.run $ do
let len = half (I# (sizeofByteArray# inner))
m <- TA.new len
TA.copyI m 0 a 0 len
let go !ix = if ix < len
then do
TA.unsafeWrite m ix (charToWord16 (f (word16ToChar (TA.unsafeIndex a ix))))
go (ix + 1)
else return ()
go 0
return m