module Country
( Country
, encodeNumeric
, decodeNumeric
, encodeEnglish
, decode
, parser
, parserUtf8
, alphaTwoUpper
, alphaThreeUpper
, alphaThreeLower
, alphaTwoLower
, decodeAlphaTwo
, decodeAlphaThree
) where
import Country.Unsafe (Country(..))
import Country.Unexposed.Encode.English (countryNameQuads)
import Country.Unexposed.Names (numberOfPossibleCodes,alphaTwoHashMap,alphaThreeHashMap,decodeMap,decodeNumeric,encodeEnglish)
import Country.Unexposed.Trie (Trie,trieFromList,trieParser)
import Country.Unexposed.TrieByte (TrieByte,trieByteFromList,trieByteParser)
import Data.Text (Text)
import Data.Word (Word16)
import Data.Primitive (indexArray,writeByteArray,indexByteArray,unsafeFreezeByteArray,newByteArray)
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 Data.Coerce (coerce)
import qualified Data.HashMap.Strict as HM
import qualified Data.Text.Array as TA
import qualified Data.Text.Encoding as TE
import qualified Data.Text.Internal as TI
import qualified Data.Attoparsec.Text as AT
import qualified Data.Attoparsec.ByteString as AB
encodeNumeric :: Country -> Word16
encodeNumeric (Country n) = 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
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
parser :: AT.Parser Country
parser = coerce (trieParser decodeTrie)
parserUtf8 :: AB.Parser Country
parserUtf8 = coerce (trieByteParser decodeTrieUtf8)
word16ToInt :: Word16 -> Int
word16ToInt = fromIntegral
charToWord16 :: Char -> Word16
charToWord16 = fromIntegral . ord
word16ToChar :: Word16 -> Char
word16ToChar = chr . fromIntegral
numberOfCountries :: Int
numberOfCountries = length countryNameQuads
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
decodeTrie :: Trie
decodeTrie = trieFromList (map (\(a,Country x) -> (a,x)) (HM.toList decodeMap))
decodeTrieUtf8 :: TrieByte
decodeTrieUtf8 = trieByteFromList (map (\(a,Country x) -> (TE.encodeUtf8 a,x)) (HM.toList decodeMap))