TrieByte.hs

module Country.Unexposed.TrieByte
  ( TrieByte
  , trieByteFromList
  , trieByteParser
  ) where

import Control.Applicative ((<|>))
import qualified Data.Attoparsec.ByteString as AB
import Data.ByteString (ByteString)
import qualified Data.ByteString as B
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HM
import qualified Data.Semigroup as SG
import Data.Word (Word16, Word8)

{- | If the value is not the max Word16 (65535), there
  is a match. This means that 65535 cannot be used, which
  is fine for this since 65535 is not used as a country code.
-}
data TrieByte = TrieByte
  { _trieValue :: {-# UNPACK #-} !Word16
  , _trieChildren :: !(HashMap Word8 TrieByte)
  }

empty :: TrieByte
empty = TrieByte placeholder HM.empty

append :: TrieByte -> TrieByte -> TrieByte
append (TrieByte v1 c1) (TrieByte v2 c2) = TrieByte (min v1 v2) (HM.unionWith append c1 c2)

placeholder :: Word16
placeholder = 0xFFFF

singleton :: ByteString -> Word16 -> TrieByte
singleton fullName code = go fullName
 where
  go :: ByteString -> TrieByte
  go name = case B.uncons name of
    Just (char, nameNext) -> TrieByte placeholder (HM.singleton char (go nameNext))
    Nothing -> TrieByte code HM.empty

instance Semigroup TrieByte where
  (<>) = append

instance Monoid TrieByte where
  mempty = empty
  mappend = (SG.<>)

trieByteFromList :: [(ByteString, Word16)] -> TrieByte
trieByteFromList = foldMap (uncurry singleton)

-- it seems like attoparsec should have some kind of convenience
-- for being able to commit to consuming a certain amount of
-- input once your certain that it will be consumed, but I cannot
-- find a way to use the api to do this.
trieByteParser :: TrieByte -> AB.Parser Word16
trieByteParser = go
 where
  go :: TrieByte -> AB.Parser Word16
  go (TrieByte value children) = do
    let keepGoing = do
          c <- AB.anyWord8
          case HM.lookup c children of
            Nothing -> fail "did not recognize country name"
            Just trieNext -> go trieNext
    if value == placeholder
      then keepGoing
      else keepGoing <|> return value