    {-# LANGUAGE CPP, BangPatterns, PatternGuards #-}
    {-# LANGUAGE DeriveDataTypeable, ScopedTypeVariables #-}
    {-# LANGUAGE MagicHash #-}
    {-# LANGUAGE UnboxedTuples #-}
    {-# OPTIONS_HADDOCK hide #-}

    module Codec.Archive.Tar.Index.IntTrie (

      IntTrie(..),
      construct,
      toList,

      IntTrieBuilder(..),
      empty,
      insert,
      finalise,
      unfinalise,

      lookup,
      TrieLookup(..),

      serialise,
      serialiseSize,
      deserialise,

      TrieNode(..),
      Completions,
      inserts,
      completionsFrom,
      flattenTrie,
      tagLeaf,
      tagNode,

      Key(..),
      Value(..),
      ) where

    import Prelude hiding (lookup)

    import Data.Typeable (Typeable)

    import qualified Data.Array.Unboxed as A
    import Data.Array.IArray  ((!))
    import qualified Data.Bits as Bits
    import Data.Word (Word32)
    import Data.Bits
    import Data.Monoid (Monoid(..))
    import Data.Monoid ((<>))
    import qualified Data.ByteString        as BS
    import qualified Data.ByteString.Lazy   as LBS
    import qualified Data.ByteString.Unsafe as BS
    import Data.ByteString.Builder          as BS
    import Control.Exception (assert)
    import qualified Data.Map.Strict        as Map
    import qualified Data.IntMap.Strict     as IntMap
    import Data.IntMap.Strict (IntMap)

    import Data.List hiding (lookup, insert)
    import Data.Function (on)
    import GHC.IO

    import Codec.Archive.Tar.Index.Utils

    -- | A compact mapping from sequences of nats to nats.
    --
    -- NOTE: The tries in this module have values /only/ at the leaves (which
    -- correspond to files), they do not have values at the branch points (which
    -- correspond to directories).
    newtype IntTrie = IntTrie (A.UArray Word32 Word32)
        deriving (IntTrie -> IntTrie -> Bool
    (IntTrie -> IntTrie -> Bool)
    -> (IntTrie -> IntTrie -> Bool) -> Eq IntTrie
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: IntTrie -> IntTrie -> Bool
    == :: IntTrie -> IntTrie -> Bool
    $c/= :: IntTrie -> IntTrie -> Bool
    /= :: IntTrie -> IntTrie -> Bool
    Eq, Int -> IntTrie -> ShowS
    [IntTrie] -> ShowS
    IntTrie -> String
    (Int -> IntTrie -> ShowS)
    -> (IntTrie -> String) -> ([IntTrie] -> ShowS) -> Show IntTrie
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> IntTrie -> ShowS
    showsPrec :: Int -> IntTrie -> ShowS
    $cshow :: IntTrie -> String
    show :: IntTrie -> String
    $cshowList :: [IntTrie] -> ShowS
    showList :: [IntTrie] -> ShowS
    Show, Typeable)

    -- | The most significant bit is used for tagging,
    -- see 'tagLeaf' / 'tagNode' below, so morally it's Word31 only.
    newtype Key = Key { Key -> Word32
    unKey :: Word32 }
      deriving (Key -> Key -> Bool
    (Key -> Key -> Bool) -> (Key -> Key -> Bool) -> Eq Key
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: Key -> Key -> Bool
    == :: Key -> Key -> Bool
    $c/= :: Key -> Key -> Bool
    /= :: Key -> Key -> Bool
    Eq, Eq Key
    Eq Key =>
    (Key -> Key -> Ordering)
    -> (Key -> Key -> Bool)
    -> (Key -> Key -> Bool)
    -> (Key -> Key -> Bool)
    -> (Key -> Key -> Bool)
    -> (Key -> Key -> Key)
    -> (Key -> Key -> Key)
    -> Ord Key
    Key -> Key -> Bool
    Key -> Key -> Ordering
    Key -> Key -> Key
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    $ccompare :: Key -> Key -> Ordering
    compare :: Key -> Key -> Ordering
    $c< :: Key -> Key -> Bool
    < :: Key -> Key -> Bool
    $c<= :: Key -> Key -> Bool
    <= :: Key -> Key -> Bool
    $c> :: Key -> Key -> Bool
    > :: Key -> Key -> Bool
    $c>= :: Key -> Key -> Bool
    >= :: Key -> Key -> Bool
    $cmax :: Key -> Key -> Key
    max :: Key -> Key -> Key
    $cmin :: Key -> Key -> Key
    min :: Key -> Key -> Key
    Ord, Int -> Key -> ShowS
    [Key] -> ShowS
    Key -> String
    (Int -> Key -> ShowS)
    -> (Key -> String) -> ([Key] -> ShowS) -> Show Key
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> Key -> ShowS
    showsPrec :: Int -> Key -> ShowS
    $cshow :: Key -> String
    show :: Key -> String
    $cshowList :: [Key] -> ShowS
    showList :: [Key] -> ShowS
    Show)

    newtype Value = Value { Value -> Word32
    unValue :: Word32 }
      deriving (Value -> Value -> Bool
    (Value -> Value -> Bool) -> (Value -> Value -> Bool) -> Eq Value
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: Value -> Value -> Bool
    == :: Value -> Value -> Bool
    $c/= :: Value -> Value -> Bool
    /= :: Value -> Value -> Bool
    Eq, Eq Value
    Eq Value =>
    (Value -> Value -> Ordering)
    -> (Value -> Value -> Bool)
    -> (Value -> Value -> Bool)
    -> (Value -> Value -> Bool)
    -> (Value -> Value -> Bool)
    -> (Value -> Value -> Value)
    -> (Value -> Value -> Value)
    -> Ord Value
    Value -> Value -> Bool
    Value -> Value -> Ordering
    Value -> Value -> Value
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    $ccompare :: Value -> Value -> Ordering
    compare :: Value -> Value -> Ordering
    $c< :: Value -> Value -> Bool
    < :: Value -> Value -> Bool
    $c<= :: Value -> Value -> Bool
    <= :: Value -> Value -> Bool
    $c> :: Value -> Value -> Bool
    > :: Value -> Value -> Bool
    $c>= :: Value -> Value -> Bool
    >= :: Value -> Value -> Bool
    $cmax :: Value -> Value -> Value
    max :: Value -> Value -> Value
    $cmin :: Value -> Value -> Value
    min :: Value -> Value -> Value
    Ord, Int -> Value -> ShowS
    [Value] -> ShowS
    Value -> String
    (Int -> Value -> ShowS)
    -> (Value -> String) -> ([Value] -> ShowS) -> Show Value
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> Value -> ShowS
    showsPrec :: Int -> Value -> ShowS
    $cshow :: Value -> String
    show :: Value -> String
    $cshowList :: [Value] -> ShowS
    showList :: [Value] -> ShowS
    Show)

    -- Compact, read-only implementation of a trie. It's intended for use with file
    -- paths, but we do that via string ids.

    -- Each node has a size and a sequence of keys followed by an equal length
    -- sequence of corresponding entries. Since we're going to flatten this into
    -- a single array then we will need to replace the trie structure with pointers
    -- represented as array offsets.

    -- Each node is a pair of arrays, one of keys and one of Either value pointer.
    -- We need to distinguish values from internal pointers. We use a tag bit:
    --
    tagLeaf, tagNode, untag :: Word32 -> Word32
    tagLeaf :: Word32 -> Word32
    tagLeaf = Word32 -> Word32
    forall a. a -> a
    id
    tagNode :: Word32 -> Word32
    tagNode = (Word32 -> Int -> Word32) -> Int -> Word32 -> Word32
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip Word32 -> Int -> Word32
    forall a. Bits a => a -> Int -> a
    Bits.setBit   Int
    31
    untag :: Word32 -> Word32
    untag   = (Word32 -> Int -> Word32) -> Int -> Word32 -> Word32
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip Word32 -> Int -> Word32
    forall a. Bits a => a -> Int -> a
    Bits.clearBit Int
    31

    isNode :: Word32 -> Bool
    isNode :: Word32 -> Bool
    isNode = (Word32 -> Int -> Bool) -> Int -> Word32 -> Bool
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip Word32 -> Int -> Bool
    forall a. Bits a => a -> Int -> Bool
    Bits.testBit Int
    31

    -------------------------------------
    -- Decoding the trie array form
    --

    completionsFrom :: IntTrie -> Word32 -> Completions
    completionsFrom :: IntTrie -> Word32 -> Completions
    completionsFrom trie :: IntTrie
    trie@(IntTrie UArray Word32 Word32
    arr) Word32
    nodeOff =
        [ (Word32 -> Key
    Key (Word32 -> Word32
    untag Word32
    key), TrieLookup
    next)
        | Word32
    keyOff <- [Word32
    keysStart..Word32
    keysEnd]
        , let key :: Word32
    key   = UArray Word32 Word32
    arr UArray Word32 Word32 -> Word32 -> Word32
    forall (a :: * -> * -> *) e i.
    (IArray a e, Ix i) =>
    a i e -> i -> e
    ! Word32
    keyOff
              entry :: Word32
    entry = UArray Word32 Word32
    arr UArray Word32 Word32 -> Word32 -> Word32
    forall (a :: * -> * -> *) e i.
    (IArray a e, Ix i) =>
    a i e -> i -> e
    ! (Word32
    keyOff Word32 -> Word32 -> Word32
    forall a. Num a => a -> a -> a
    + Word32
    nodeSize)
              next :: TrieLookup
    next | Word32 -> Bool
    isNode Word32
    key = Completions -> TrieLookup
    Completions (IntTrie -> Word32 -> Completions
    completionsFrom IntTrie
    trie Word32
    entry)
                   | Bool
    otherwise  = Value -> TrieLookup
    Entry (Word32 -> Value
    Value Word32
    entry)
        ]
      where
        nodeSize :: Word32
    nodeSize  = UArray Word32 Word32
    arr UArray Word32 Word32 -> Word32 -> Word32
    forall (a :: * -> * -> *) e i.
    (IArray a e, Ix i) =>
    a i e -> i -> e
    ! Word32
    nodeOff
        keysStart :: Word32
    keysStart = Word32
    nodeOff Word32 -> Word32 -> Word32
    forall a. Num a => a -> a -> a
    + Word32
    1
        keysEnd :: Word32
    keysEnd   = Word32
    nodeOff Word32 -> Word32 -> Word32
    forall a. Num a => a -> a -> a
    + Word32
    nodeSize

    -- | Convert the trie to a list
    --
    -- This is the left inverse to 'construct' (modulo ordering).
    toList :: IntTrie -> [([Key], Value)]
    toList :: IntTrie -> [([Key], Value)]
    toList = ((Key, TrieLookup) -> [([Key], Value)])
    -> Completions -> [([Key], Value)]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap ([Key] -> (Key, TrieLookup) -> [([Key], Value)]
    aux []) (Completions -> [([Key], Value)])
    -> (IntTrie -> Completions) -> IntTrie -> [([Key], Value)]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (IntTrie -> Word32 -> Completions
    `completionsFrom` Word32
    0)
      where
        aux :: [Key] -> (Key, TrieLookup) -> [([Key], Value)]
        aux :: [Key] -> (Key, TrieLookup) -> [([Key], Value)]
    aux [Key]
    ks (Key
    k, Entry Value
    v)        = [([Key] -> [Key]
    forall a. [a] -> [a]
    reverse (Key
    kKey -> [Key] -> [Key]
    forall a. a -> [a] -> [a]
    :[Key]
    ks), Value
    v)]
        aux [Key]
    ks (Key
    k, Completions Completions
    cs) = ((Key, TrieLookup) -> [([Key], Value)])
    -> Completions -> [([Key], Value)]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap ([Key] -> (Key, TrieLookup) -> [([Key], Value)]
    aux (Key
    kKey -> [Key] -> [Key]
    forall a. a -> [a] -> [a]
    :[Key]
    ks)) Completions
    cs

    -------------------------------------
    -- Toplevel trie array construction
    --

    -- So constructing the t'IntTrie' as a whole is just a matter of stringing
    -- together all the bits

    -- | Build an t'IntTrie' from a bunch of (key, value) pairs, where the keys
    -- are sequences.
    --
    construct :: [([Key], Value)] -> IntTrie
    construct :: [([Key], Value)] -> IntTrie
    construct = IntTrieBuilder -> IntTrie
    finalise (IntTrieBuilder -> IntTrie)
    -> ([([Key], Value)] -> IntTrieBuilder)
    -> [([Key], Value)]
    -> IntTrie
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ([([Key], Value)] -> IntTrieBuilder -> IntTrieBuilder)
    -> IntTrieBuilder -> [([Key], Value)] -> IntTrieBuilder
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip [([Key], Value)] -> IntTrieBuilder -> IntTrieBuilder
    inserts IntTrieBuilder
    empty

    ---------------------------------
    -- Looking up in the trie array
    --

    data TrieLookup = Entry !Value | Completions Completions
      deriving (TrieLookup -> TrieLookup -> Bool
    (TrieLookup -> TrieLookup -> Bool)
    -> (TrieLookup -> TrieLookup -> Bool) -> Eq TrieLookup
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: TrieLookup -> TrieLookup -> Bool
    == :: TrieLookup -> TrieLookup -> Bool
    $c/= :: TrieLookup -> TrieLookup -> Bool
    /= :: TrieLookup -> TrieLookup -> Bool
    Eq, Eq TrieLookup
    Eq TrieLookup =>
    (TrieLookup -> TrieLookup -> Ordering)
    -> (TrieLookup -> TrieLookup -> Bool)
    -> (TrieLookup -> TrieLookup -> Bool)
    -> (TrieLookup -> TrieLookup -> Bool)
    -> (TrieLookup -> TrieLookup -> Bool)
    -> (TrieLookup -> TrieLookup -> TrieLookup)
    -> (TrieLookup -> TrieLookup -> TrieLookup)
    -> Ord TrieLookup
    TrieLookup -> TrieLookup -> Bool
    TrieLookup -> TrieLookup -> Ordering
    TrieLookup -> TrieLookup -> TrieLookup
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    $ccompare :: TrieLookup -> TrieLookup -> Ordering
    compare :: TrieLookup -> TrieLookup -> Ordering
    $c< :: TrieLookup -> TrieLookup -> Bool
    < :: TrieLookup -> TrieLookup -> Bool
    $c<= :: TrieLookup -> TrieLookup -> Bool
    <= :: TrieLookup -> TrieLookup -> Bool
    $c> :: TrieLookup -> TrieLookup -> Bool
    > :: TrieLookup -> TrieLookup -> Bool
    $c>= :: TrieLookup -> TrieLookup -> Bool
    >= :: TrieLookup -> TrieLookup -> Bool
    $cmax :: TrieLookup -> TrieLookup -> TrieLookup
    max :: TrieLookup -> TrieLookup -> TrieLookup
    $cmin :: TrieLookup -> TrieLookup -> TrieLookup
    min :: TrieLookup -> TrieLookup -> TrieLookup
    Ord, Int -> TrieLookup -> ShowS
    [TrieLookup] -> ShowS
    TrieLookup -> String
    (Int -> TrieLookup -> ShowS)
    -> (TrieLookup -> String)
    -> ([TrieLookup] -> ShowS)
    -> Show TrieLookup
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> TrieLookup -> ShowS
    showsPrec :: Int -> TrieLookup -> ShowS
    $cshow :: TrieLookup -> String
    show :: TrieLookup -> String
    $cshowList :: [TrieLookup] -> ShowS
    showList :: [TrieLookup] -> ShowS
    Show)

    type Completions = [(Key, TrieLookup)]

    lookup :: IntTrie -> [Key] -> Maybe TrieLookup
    lookup :: IntTrie -> [Key] -> Maybe TrieLookup
    lookup trie :: IntTrie
    trie@(IntTrie UArray Word32 Word32
    arr) = Word32 -> [Key] -> Maybe TrieLookup
    go Word32
    0
      where
        go :: Word32 -> [Key] -> Maybe TrieLookup
        go :: Word32 -> [Key] -> Maybe TrieLookup
    go Word32
    nodeOff []     = TrieLookup -> Maybe TrieLookup
    forall a. a -> Maybe a
    Just (Word32 -> TrieLookup
    completions Word32
    nodeOff)
        go Word32
    nodeOff (Key
    k:[Key]
    ks) = case Word32 -> Word32 -> Maybe Word32
    search Word32
    nodeOff (Word32 -> Word32
    tagLeaf Word32
    k') of
          Just Word32
    entryOff
            | [Key] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [Key]
    ks   -> TrieLookup -> Maybe TrieLookup
    forall a. a -> Maybe a
    Just (Word32 -> TrieLookup
    entry Word32
    entryOff)
            | Bool
    otherwise -> Maybe TrieLookup
    forall a. Maybe a
    Nothing
          Maybe Word32
    Nothing       -> case Word32 -> Word32 -> Maybe Word32
    search Word32
    nodeOff (Word32 -> Word32
    tagNode Word32
    k') of
            Maybe Word32
    Nothing       -> Maybe TrieLookup
    forall a. Maybe a
    Nothing
            Just Word32
    entryOff -> Word32 -> [Key] -> Maybe TrieLookup
    go (UArray Word32 Word32
    arr UArray Word32 Word32 -> Word32 -> Word32
    forall (a :: * -> * -> *) e i.
    (IArray a e, Ix i) =>
    a i e -> i -> e
    ! Word32
    entryOff) [Key]
    ks
          where
            k' :: Word32
    k' = Key -> Word32
    unKey Key
    k

        entry :: Word32 -> TrieLookup
    entry       Word32
    entryOff = Value -> TrieLookup
    Entry (Word32 -> Value
    Value (UArray Word32 Word32
    arr UArray Word32 Word32 -> Word32 -> Word32
    forall (a :: * -> * -> *) e i.
    (IArray a e, Ix i) =>
    a i e -> i -> e
    ! Word32
    entryOff))
        completions :: Word32 -> TrieLookup
    completions Word32
    nodeOff  = Completions -> TrieLookup
    Completions (IntTrie -> Word32 -> Completions
    completionsFrom IntTrie
    trie Word32
    nodeOff)

        search :: Word32 -> Word32 -> Maybe Word32
        search :: Word32 -> Word32 -> Maybe Word32
    search Word32
    nodeOff Word32
    key = (Word32 -> Word32) -> Maybe Word32 -> Maybe Word32
    forall a b. (a -> b) -> Maybe a -> Maybe b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap (Word32 -> Word32 -> Word32
    forall a. Num a => a -> a -> a
    +Word32
    nodeSize) (Word32 -> Word32 -> Word32 -> Maybe Word32
    bsearch Word32
    keysStart Word32
    keysEnd Word32
    key)
          where
            nodeSize :: Word32
    nodeSize  = UArray Word32 Word32
    arr UArray Word32 Word32 -> Word32 -> Word32
    forall (a :: * -> * -> *) e i.
    (IArray a e, Ix i) =>
    a i e -> i -> e
    ! Word32
    nodeOff
            keysStart :: Word32
    keysStart = Word32
    nodeOff Word32 -> Word32 -> Word32
    forall a. Num a => a -> a -> a
    + Word32
    1
            keysEnd :: Word32
    keysEnd   = Word32
    nodeOff Word32 -> Word32 -> Word32
    forall a. Num a => a -> a -> a
    + Word32
    nodeSize

        bsearch :: Word32 -> Word32 -> Word32 -> Maybe Word32
        bsearch :: Word32 -> Word32 -> Word32 -> Maybe Word32
    bsearch Word32
    a Word32
    b Word32
    key
          | Word32
    a Word32 -> Word32 -> Bool
    forall a. Ord a => a -> a -> Bool
    > Word32
    b     = Maybe Word32
    forall a. Maybe a
    Nothing
          | Bool
    otherwise = case Word32 -> Word32 -> Ordering
    forall a. Ord a => a -> a -> Ordering
    compare Word32
    key (UArray Word32 Word32
    arr UArray Word32 Word32 -> Word32 -> Word32
    forall (a :: * -> * -> *) e i.
    (IArray a e, Ix i) =>
    a i e -> i -> e
    ! Word32
    mid) of
              Ordering
    LT -> Word32 -> Word32 -> Word32 -> Maybe Word32
    bsearch Word32
    a (Word32
    midWord32 -> Word32 -> Word32
    forall a. Num a => a -> a -> a
    -Word32
    1) Word32
    key
              Ordering
    EQ -> Word32 -> Maybe Word32
    forall a. a -> Maybe a
    Just Word32
    mid
              Ordering
    GT -> Word32 -> Word32 -> Word32 -> Maybe Word32
    bsearch (Word32
    midWord32 -> Word32 -> Word32
    forall a. Num a => a -> a -> a
    +Word32
    1) Word32
    b Word32
    key
          where mid :: Word32
    mid = (Word32
    a Word32 -> Word32 -> Word32
    forall a. Num a => a -> a -> a
    + Word32
    b) Word32 -> Word32 -> Word32
    forall a. Integral a => a -> a -> a
    `div` Word32
    2

    -------------------------
    -- Building Tries
    --

    newtype IntTrieBuilder = IntTrieBuilder (IntMap TrieNode)
      deriving (Int -> IntTrieBuilder -> ShowS
    [IntTrieBuilder] -> ShowS
    IntTrieBuilder -> String
    (Int -> IntTrieBuilder -> ShowS)
    -> (IntTrieBuilder -> String)
    -> ([IntTrieBuilder] -> ShowS)
    -> Show IntTrieBuilder
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> IntTrieBuilder -> ShowS
    showsPrec :: Int -> IntTrieBuilder -> ShowS
    $cshow :: IntTrieBuilder -> String
    show :: IntTrieBuilder -> String
    $cshowList :: [IntTrieBuilder] -> ShowS
    showList :: [IntTrieBuilder] -> ShowS
    Show, IntTrieBuilder -> IntTrieBuilder -> Bool
    (IntTrieBuilder -> IntTrieBuilder -> Bool)
    -> (IntTrieBuilder -> IntTrieBuilder -> Bool) -> Eq IntTrieBuilder
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: IntTrieBuilder -> IntTrieBuilder -> Bool
    == :: IntTrieBuilder -> IntTrieBuilder -> Bool
    $c/= :: IntTrieBuilder -> IntTrieBuilder -> Bool
    /= :: IntTrieBuilder -> IntTrieBuilder -> Bool
    Eq)

    data TrieNode = TrieLeaf {-# UNPACK #-} !Word32
                  | TrieNode !IntTrieBuilder
      deriving (Int -> TrieNode -> ShowS
    [TrieNode] -> ShowS
    TrieNode -> String
    (Int -> TrieNode -> ShowS)
    -> (TrieNode -> String) -> ([TrieNode] -> ShowS) -> Show TrieNode
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: Int -> TrieNode -> ShowS
    showsPrec :: Int -> TrieNode -> ShowS
    $cshow :: TrieNode -> String
    show :: TrieNode -> String
    $cshowList :: [TrieNode] -> ShowS
    showList :: [TrieNode] -> ShowS
    Show, TrieNode -> TrieNode -> Bool
    (TrieNode -> TrieNode -> Bool)
    -> (TrieNode -> TrieNode -> Bool) -> Eq TrieNode
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: TrieNode -> TrieNode -> Bool
    == :: TrieNode -> TrieNode -> Bool
    $c/= :: TrieNode -> TrieNode -> Bool
    /= :: TrieNode -> TrieNode -> Bool
    Eq)

    empty :: IntTrieBuilder
    empty :: IntTrieBuilder
    empty = IntMap TrieNode -> IntTrieBuilder
    IntTrieBuilder IntMap TrieNode
    forall a. IntMap a
    IntMap.empty

    insert :: [Key] -> Value
           -> IntTrieBuilder -> IntTrieBuilder
    insert :: [Key] -> Value -> IntTrieBuilder -> IntTrieBuilder
    insert []    Value
    _v IntTrieBuilder
    t = IntTrieBuilder
    t
    insert (Key
    k:[Key]
    ks) Value
    v IntTrieBuilder
    t = Int -> [Int] -> Word32 -> IntTrieBuilder -> IntTrieBuilder
    insertTrie
      (Word32 -> Int
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral (Key -> Word32
    unKey Key
    k) :: Int)
      ((Key -> Int) -> [Key] -> [Int]
    forall a b. (a -> b) -> [a] -> [b]
    map (Word32 -> Int
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral (Word32 -> Int) -> (Key -> Word32) -> Key -> Int
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Key -> Word32
    unKey) [Key]
    ks :: [Int])
      (Value -> Word32
    unValue Value
    v)
      IntTrieBuilder
    t

    insertTrie :: Int -> [Int] -> Word32
               -> IntTrieBuilder -> IntTrieBuilder
    insertTrie :: Int -> [Int] -> Word32 -> IntTrieBuilder -> IntTrieBuilder
    insertTrie Int
    k [Int]
    ks Word32
    v (IntTrieBuilder IntMap TrieNode
    t) =
        IntMap TrieNode -> IntTrieBuilder
    IntTrieBuilder (IntMap TrieNode -> IntTrieBuilder)
    -> IntMap TrieNode -> IntTrieBuilder
    forall a b. (a -> b) -> a -> b
    $
          (Maybe TrieNode -> Maybe TrieNode)
    -> Int -> IntMap TrieNode -> IntMap TrieNode
    forall a. (Maybe a -> Maybe a) -> Int -> IntMap a -> IntMap a
    IntMap.alter (\Maybe TrieNode
    t' -> TrieNode -> Maybe TrieNode
    forall a. a -> Maybe a
    Just (TrieNode -> Maybe TrieNode) -> TrieNode -> Maybe TrieNode
    forall a b. (a -> b) -> a -> b
    $! TrieNode -> (TrieNode -> TrieNode) -> Maybe TrieNode -> TrieNode
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe ([Int] -> Word32 -> TrieNode
    freshTrieNode  [Int]
    ks Word32
    v)
                                             ([Int] -> Word32 -> TrieNode -> TrieNode
    insertTrieNode [Int]
    ks Word32
    v) Maybe TrieNode
    t')
                       Int
    k IntMap TrieNode
    t

    insertTrieNode :: [Int] -> Word32 -> TrieNode -> TrieNode
    insertTrieNode :: [Int] -> Word32 -> TrieNode -> TrieNode
    insertTrieNode []     Word32
    v  TrieNode
    _           = Word32 -> TrieNode
    TrieLeaf Word32
    v
    insertTrieNode (Int
    k:[Int]
    ks) Word32
    v (TrieLeaf Word32
    _) = IntTrieBuilder -> TrieNode
    TrieNode (Int -> [Int] -> Word32 -> IntTrieBuilder
    freshTrie  Int
    k [Int]
    ks Word32
    v)
    insertTrieNode (Int
    k:[Int]
    ks) Word32
    v (TrieNode IntTrieBuilder
    t) = IntTrieBuilder -> TrieNode
    TrieNode (Int -> [Int] -> Word32 -> IntTrieBuilder -> IntTrieBuilder
    insertTrie Int
    k [Int]
    ks Word32
    v IntTrieBuilder
    t)

    freshTrie :: Int -> [Int] -> Word32 -> IntTrieBuilder
    freshTrie :: Int -> [Int] -> Word32 -> IntTrieBuilder
    freshTrie Int
    k []      Word32
    v =
      IntMap TrieNode -> IntTrieBuilder
    IntTrieBuilder (Int -> TrieNode -> IntMap TrieNode
    forall a. Int -> a -> IntMap a
    IntMap.singleton Int
    k (Word32 -> TrieNode
    TrieLeaf Word32
    v))
    freshTrie Int
    k (Int
    k':[Int]
    ks) Word32
    v =
      IntMap TrieNode -> IntTrieBuilder
    IntTrieBuilder (Int -> TrieNode -> IntMap TrieNode
    forall a. Int -> a -> IntMap a
    IntMap.singleton Int
    k (IntTrieBuilder -> TrieNode
    TrieNode (Int -> [Int] -> Word32 -> IntTrieBuilder
    freshTrie Int
    k' [Int]
    ks Word32
    v)))

    freshTrieNode :: [Int] -> Word32 -> TrieNode
    freshTrieNode :: [Int] -> Word32 -> TrieNode
    freshTrieNode []     Word32
    v = Word32 -> TrieNode
    TrieLeaf Word32
    v
    freshTrieNode (Int
    k:[Int]
    ks) Word32
    v = IntTrieBuilder -> TrieNode
    TrieNode (Int -> [Int] -> Word32 -> IntTrieBuilder
    freshTrie Int
    k [Int]
    ks Word32
    v)

    inserts :: [([Key], Value)]
            -> IntTrieBuilder -> IntTrieBuilder
    inserts :: [([Key], Value)] -> IntTrieBuilder -> IntTrieBuilder
    inserts [([Key], Value)]
    kvs IntTrieBuilder
    t = (IntTrieBuilder -> ([Key], Value) -> IntTrieBuilder)
    -> IntTrieBuilder -> [([Key], Value)] -> IntTrieBuilder
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl' (\IntTrieBuilder
    t' ([Key]
    ks, Value
    v) -> [Key] -> Value -> IntTrieBuilder -> IntTrieBuilder
    insert [Key]
    ks Value
    v IntTrieBuilder
    t') IntTrieBuilder
    t [([Key], Value)]
    kvs

    finalise :: IntTrieBuilder -> IntTrie
    finalise :: IntTrieBuilder -> IntTrie
    finalise IntTrieBuilder
    trie =
        UArray Word32 Word32 -> IntTrie
    IntTrie (UArray Word32 Word32 -> IntTrie)
    -> UArray Word32 Word32 -> IntTrie
    forall a b. (a -> b) -> a -> b
    $
          (Word32, Word32) -> [Word32] -> UArray Word32 Word32
    forall (a :: * -> * -> *) e i.
    (IArray a e, Ix i) =>
    (i, i) -> [e] -> a i e
    A.listArray (Word32
    0, Int -> Word32
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral (IntTrieBuilder -> Int
    flatTrieLength IntTrieBuilder
    trie) Word32 -> Word32 -> Word32
    forall a. Num a => a -> a -> a
    - Word32
    1)
                      (IntTrieBuilder -> [Word32]
    flattenTrie IntTrieBuilder
    trie)

    unfinalise :: IntTrie -> IntTrieBuilder
    unfinalise :: IntTrie -> IntTrieBuilder
    unfinalise IntTrie
    trie =
        Completions -> IntTrieBuilder
    go (IntTrie -> Word32 -> Completions
    completionsFrom IntTrie
    trie Word32
    0)
      where
        go :: Completions -> IntTrieBuilder
    go Completions
    kns =
          IntMap TrieNode -> IntTrieBuilder
    IntTrieBuilder (IntMap TrieNode -> IntTrieBuilder)
    -> IntMap TrieNode -> IntTrieBuilder
    forall a b. (a -> b) -> a -> b
    $
            [(Int, TrieNode)] -> IntMap TrieNode
    forall a. [(Int, a)] -> IntMap a
    IntMap.fromList
              [ (Word32 -> Int
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral (Key -> Word32
    unKey Key
    k) :: Int, TrieNode
    t)
              | (Key
    k, TrieLookup
    n) <- Completions
    kns
              , let t :: TrieNode
    t = case TrieLookup
    n of
                          Entry       Value
    v    -> Word32 -> TrieNode
    TrieLeaf (Value -> Word32
    unValue Value
    v)
                          Completions Completions
    kns' -> IntTrieBuilder -> TrieNode
    TrieNode (Completions -> IntTrieBuilder
    go Completions
    kns')
              ]

    ---------------------------------
    -- Flattening Tries
    --

    type Offset = Int

    flatTrieLength :: IntTrieBuilder -> Int
    flatTrieLength :: IntTrieBuilder -> Int
    flatTrieLength (IntTrieBuilder IntMap TrieNode
    tns) =
        Int
    1
      Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    2 Int -> Int -> Int
    forall a. Num a => a -> a -> a
    * IntMap TrieNode -> Int
    forall a. IntMap a -> Int
    IntMap.size IntMap TrieNode
    tns
      Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + [Int] -> Int
    forall a. Num a => [a] -> a
    forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
    sum [ IntTrieBuilder -> Int
    flatTrieLength IntTrieBuilder
    n | TrieNode IntTrieBuilder
    n <- IntMap TrieNode -> [TrieNode]
    forall a. IntMap a -> [a]
    IntMap.elems IntMap TrieNode
    tns ]

    -- This is a breadth-first traversal. We keep a list of the tries that we are
    -- to write out next. Each of these have an offset allocated to them at the
    -- time we put them into the list. We keep a running offset so we know where
    -- to allocate next.
    --
    flattenTrie :: IntTrieBuilder -> [Word32]
    flattenTrie :: IntTrieBuilder -> [Word32]
    flattenTrie IntTrieBuilder
    trie = Q IntTrieBuilder -> Int -> [Word32]
    go ([IntTrieBuilder] -> Q IntTrieBuilder
    forall a. [a] -> Q a
    queue [IntTrieBuilder
    trie]) (IntTrieBuilder -> Int
    size IntTrieBuilder
    trie)
      where
        size :: IntTrieBuilder -> Int
    size (IntTrieBuilder IntMap TrieNode
    tns) = Int
    1 Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    2 Int -> Int -> Int
    forall a. Num a => a -> a -> a
    * IntMap TrieNode -> Int
    forall a. IntMap a -> Int
    IntMap.size IntMap TrieNode
    tns

        go :: Q IntTrieBuilder -> Offset -> [Word32]
        go :: Q IntTrieBuilder -> Int -> [Word32]
    go Q IntTrieBuilder
    todo !Int
    offset =
          case Q IntTrieBuilder -> Maybe (IntTrieBuilder, Q IntTrieBuilder)
    forall a. Q a -> Maybe (a, Q a)
    dequeue Q IntTrieBuilder
    todo of
            Maybe (IntTrieBuilder, Q IntTrieBuilder)
    Nothing                   -> []
            Just (IntTrieBuilder IntMap TrieNode
    tnodes, Q IntTrieBuilder
    tries) ->
                [Word32]
    flat [Word32] -> [Word32] -> [Word32]
    forall a. [a] -> [a] -> [a]
    ++ Q IntTrieBuilder -> Int -> [Word32]
    go Q IntTrieBuilder
    tries' Int
    offset'
              where
                !count :: Int
    count = IntMap TrieNode -> Int
    forall a. IntMap a -> Int
    IntMap.size IntMap TrieNode
    tnodes
                flat :: [Word32]
    flat   = Int -> Word32
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Int
    count
                       Word32 -> [Word32] -> [Word32]
    forall a. a -> [a] -> [a]
    : Map Word32 Word32 -> [Word32]
    forall k a. Map k a -> [k]
    Map.keys  Map Word32 Word32
    keysValues
                      [Word32] -> [Word32] -> [Word32]
    forall a. [a] -> [a] -> [a]
    ++ Map Word32 Word32 -> [Word32]
    forall k a. Map k a -> [a]
    Map.elems Map Word32 Word32
    keysValues
                (!Int
    offset', !Map Word32 Word32
    keysValues, !Q IntTrieBuilder
    tries') =
                  ((Int, Map Word32 Word32, Q IntTrieBuilder)
     -> Int -> TrieNode -> (Int, Map Word32 Word32, Q IntTrieBuilder))
    -> (Int, Map Word32 Word32, Q IntTrieBuilder)
    -> IntMap TrieNode
    -> (Int, Map Word32 Word32, Q IntTrieBuilder)
    forall a b. (a -> Int -> b -> a) -> a -> IntMap b -> a
    IntMap.foldlWithKey' (Int, Map Word32 Word32, Q IntTrieBuilder)
    -> Int -> TrieNode -> (Int, Map Word32 Word32, Q IntTrieBuilder)
    accumNodes
                                       (Int
    offset, Map Word32 Word32
    forall k a. Map k a
    Map.empty, Q IntTrieBuilder
    tries)
                                       IntMap TrieNode
    tnodes

        accumNodes :: (Offset, Map.Map Word32 Word32, Q IntTrieBuilder)
                   -> Int -> TrieNode
                   -> (Offset, Map.Map Word32 Word32, Q IntTrieBuilder)
        accumNodes :: (Int, Map Word32 Word32, Q IntTrieBuilder)
    -> Int -> TrieNode -> (Int, Map Word32 Word32, Q IntTrieBuilder)
    accumNodes (!Int
    off, !Map Word32 Word32
    kvs, !Q IntTrieBuilder
    tries) !Int
    k (TrieLeaf Word32
    v) =
            (Int
    off, Map Word32 Word32
    kvs', Q IntTrieBuilder
    tries)
          where
            kvs' :: Map Word32 Word32
    kvs' = Word32 -> Word32 -> Map Word32 Word32 -> Map Word32 Word32
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    Map.insert (Word32 -> Word32
    tagLeaf (Int -> Word32
    int2Word32 Int
    k)) Word32
    v Map Word32 Word32
    kvs

        accumNodes (!Int
    off, !Map Word32 Word32
    kvs, !Q IntTrieBuilder
    tries) !Int
    k (TrieNode IntTrieBuilder
    t) =
            (Int
    off Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + IntTrieBuilder -> Int
    size IntTrieBuilder
    t, Map Word32 Word32
    kvs', Q IntTrieBuilder
    tries')
          where
            kvs' :: Map Word32 Word32
    kvs'   = Word32 -> Word32 -> Map Word32 Word32 -> Map Word32 Word32
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    Map.insert (Word32 -> Word32
    tagNode (Int -> Word32
    int2Word32 Int
    k)) (Int -> Word32
    int2Word32 Int
    off) Map Word32 Word32
    kvs
            tries' :: Q IntTrieBuilder
    tries' = Q IntTrieBuilder -> IntTrieBuilder -> Q IntTrieBuilder
    forall a. Q a -> a -> Q a
    enqueue Q IntTrieBuilder
    tries IntTrieBuilder
    t

    data Q a = Q [a] [a]

    queue :: [a] -> Q a
    queue :: forall a. [a] -> Q a
    queue [a]
    xs = [a] -> [a] -> Q a
    forall a. [a] -> [a] -> Q a
    Q [a]
    xs []

    enqueue :: Q a -> a -> Q a
    enqueue :: forall a. Q a -> a -> Q a
    enqueue (Q [a]
    front  [a]
    back) a
    x = [a] -> [a] -> Q a
    forall a. [a] -> [a] -> Q a
    Q [a]
    front (a
    x a -> [a] -> [a]
    forall a. a -> [a] -> [a]
    : [a]
    back)

    dequeue :: Q a -> Maybe (a, Q a)
    dequeue :: forall a. Q a -> Maybe (a, Q a)
    dequeue (Q (a
    x:[a]
    xs) [a]
    back)    = (a, Q a) -> Maybe (a, Q a)
    forall a. a -> Maybe a
    Just (a
    x, [a] -> [a] -> Q a
    forall a. [a] -> [a] -> Q a
    Q [a]
    xs [a]
    back)
    dequeue (Q []     [a]
    back)    = case [a] -> [a]
    forall a. [a] -> [a]
    reverse [a]
    back of
                                   a
    x:[a]
    xs -> (a, Q a) -> Maybe (a, Q a)
    forall a. a -> Maybe a
    Just (a
    x, [a] -> [a] -> Q a
    forall a. [a] -> [a] -> Q a
    Q [a]
    xs [])
                                   []   -> Maybe (a, Q a)
    forall a. Maybe a
    Nothing

    int2Word32 :: Int -> Word32
    int2Word32 :: Int -> Word32
    int2Word32 = Int -> Word32
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral

    -------------------------
    -- (de)serialisation
    --

    serialise :: IntTrie -> BS.Builder
    serialise :: IntTrie -> Builder
    serialise (IntTrie UArray Word32 Word32
    arr) =
        let (Word32
    _, !Word32
    ixEnd) = UArray Word32 Word32 -> (Word32, Word32)
    forall i. Ix i => UArray i Word32 -> (i, i)
    forall (a :: * -> * -> *) e i.
    (IArray a e, Ix i) =>
    a i e -> (i, i)
    A.bounds UArray Word32 Word32
    arr in
        Word32 -> Builder
    BS.word32BE (Word32
    ixEndWord32 -> Word32 -> Word32
    forall a. Num a => a -> a -> a
    +Word32
    1)
     Builder -> Builder -> Builder
    forall a. Semigroup a => a -> a -> a
    <> (Word32 -> Builder -> Builder) -> Builder -> [Word32] -> Builder
    forall a b. (a -> b -> b) -> b -> [a] -> b
    forall (t :: * -> *) a b.
    Foldable t =>
    (a -> b -> b) -> b -> t a -> b
    foldr (\Word32
    n Builder
    r -> Word32 -> Builder
    BS.word32BE Word32
    n Builder -> Builder -> Builder
    forall a. Semigroup a => a -> a -> a
    <> Builder
    r) Builder
    forall a. Monoid a => a
    mempty (UArray Word32 Word32 -> [Word32]
    forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [e]
    A.elems UArray Word32 Word32
    arr)

    serialiseSize :: IntTrie -> Int
    serialiseSize :: IntTrie -> Int
    serialiseSize (IntTrie UArray Word32 Word32
    arr) =
        let (Word32
    _, Word32
    ixEnd) = UArray Word32 Word32 -> (Word32, Word32)
    forall i. Ix i => UArray i Word32 -> (i, i)
    forall (a :: * -> * -> *) e i.
    (IArray a e, Ix i) =>
    a i e -> (i, i)
    A.bounds UArray Word32 Word32
    arr in
        Int
    4
      Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    4 Int -> Int -> Int
    forall a. Num a => a -> a -> a
    * (Word32 -> Int
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Word32
    ixEnd Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    1)

    deserialise :: BS.ByteString -> Maybe (IntTrie, BS.ByteString)
    deserialise :: ByteString -> Maybe (IntTrie, ByteString)
    deserialise ByteString
    bs
      | ByteString -> Int
    BS.length ByteString
    bs Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    >= Int
    4
      , let lenArr :: Word32
    lenArr   = ByteString -> Int -> Word32
    readWord32BE ByteString
    bs Int
    0
            lenTotal :: Int
    lenTotal = Int
    4 Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    4 Int -> Int -> Int
    forall a. Num a => a -> a -> a
    * Word32 -> Int
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Word32
    lenArr
      , ByteString -> Int
    BS.length ByteString
    bs Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    >= Int
    4 Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    4 Int -> Int -> Int
    forall a. Num a => a -> a -> a
    * Word32 -> Int
    forall a b. (Integral a, Num b) => a -> b
    fromIntegral Word32
    lenArr
      , let !bs_without_len :: ByteString
    bs_without_len = Int -> ByteString -> ByteString
    BS.unsafeDrop Int
    4 ByteString
    bs
            !bs_remaining :: ByteString
    bs_remaining = Int -> ByteString -> ByteString
    BS.unsafeDrop Int
    lenTotal ByteString
    bs
            !arr :: UArray Word32 Word32
    arr = IO (UArray Word32 Word32) -> UArray Word32 Word32
    forall a. IO a -> a
    unsafePerformIO (IO (UArray Word32 Word32) -> UArray Word32 Word32)
    -> IO (UArray Word32 Word32) -> UArray Word32 Word32
    forall a b. (a -> b) -> a -> b
    $ Word32 -> ByteString -> IO (UArray Word32 Word32)
    forall i.
    (Integral i, Num i) =>
    i -> ByteString -> IO (UArray i Word32)
    beToLe Word32
    lenArr ByteString
    bs_without_len
      = (IntTrie, ByteString) -> Maybe (IntTrie, ByteString)
    forall a. a -> Maybe a
    Just (UArray Word32 Word32 -> IntTrie
    IntTrie UArray Word32 Word32
    arr, ByteString
    bs_remaining)

      | Bool
    otherwise
      = Maybe (IntTrie, ByteString)
    forall a. Maybe a
    Nothing