{-# LANGUAGE CPP, BangPatterns, PatternGuards #-}
{-# LANGUAGE DeriveDataTypeable, ScopedTypeVariables #-}

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)

-- | 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
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: IntTrie -> IntTrie -> Bool
$c/= :: IntTrie -> IntTrie -> Bool
== :: IntTrie -> IntTrie -> Bool
$c== :: IntTrie -> IntTrie -> Bool
Eq, Int -> IntTrie -> ShowS
[IntTrie] -> ShowS
IntTrie -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [IntTrie] -> ShowS
$cshowList :: [IntTrie] -> ShowS
show :: IntTrie -> String
$cshow :: IntTrie -> String
showsPrec :: Int -> IntTrie -> ShowS
$cshowsPrec :: Int -> 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
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Key -> Key -> Bool
$c/= :: Key -> Key -> Bool
== :: Key -> Key -> Bool
$c== :: Key -> Key -> Bool
Eq, Eq 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
min :: Key -> Key -> Key
$cmin :: Key -> Key -> Key
max :: Key -> Key -> Key
$cmax :: Key -> Key -> Key
>= :: Key -> Key -> Bool
$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
compare :: Key -> Key -> Ordering
$ccompare :: Key -> Key -> Ordering
Ord, Int -> Key -> ShowS
[Key] -> ShowS
Key -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Key] -> ShowS
$cshowList :: [Key] -> ShowS
show :: Key -> String
$cshow :: Key -> String
showsPrec :: Int -> Key -> ShowS
$cshowsPrec :: Int -> Key -> ShowS
Show)

newtype Value = Value { Value -> Word32
unValue :: Word32 }
  deriving (Value -> Value -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Value -> Value -> Bool
$c/= :: Value -> Value -> Bool
== :: Value -> Value -> Bool
$c== :: Value -> Value -> Bool
Eq, Eq 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
min :: Value -> Value -> Value
$cmin :: Value -> Value -> Value
max :: Value -> Value -> Value
$cmax :: Value -> Value -> Value
>= :: Value -> Value -> Bool
$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
compare :: Value -> Value -> Ordering
$ccompare :: Value -> Value -> Ordering
Ord, Int -> Value -> ShowS
[Value] -> ShowS
Value -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Value] -> ShowS
$cshowList :: [Value] -> ShowS
show :: Value -> String
$cshow :: Value -> String
showsPrec :: Int -> Value -> ShowS
$cshowsPrec :: Int -> 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 = forall a. a -> a
id
tagNode :: Word32 -> Word32
tagNode = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Bits a => a -> Int -> a
Bits.setBit   Int
31
untag :: Word32 -> Word32
untag   = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Bits a => a -> Int -> a
Bits.clearBit Int
31

isNode :: Word32 -> Bool
isNode :: Word32 -> Bool
isNode = forall a b c. (a -> b -> c) -> b -> a -> c
flip 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 forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
! Word32
keyOff
          entry :: Word32
entry = UArray Word32 Word32
arr forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
! (Word32
keyOff 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 forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
! Word32
nodeOff
    keysStart :: Word32
keysStart = Word32
nodeOff forall a. Num a => a -> a -> a
+ Word32
1
    keysEnd :: Word32
keysEnd   = Word32
nodeOff 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 = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Key] -> (Key, TrieLookup) -> [([Key], Value)]
aux []) 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)        = [(forall a. [a] -> [a]
reverse (Key
kforall a. a -> [a] -> [a]
:[Key]
ks), Value
v)]
    aux [Key]
ks (Key
k, Completions Completions
cs) = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([Key] -> (Key, TrieLookup) -> [([Key], Value)]
aux (Key
kforall a. a -> [a] -> [a]
:[Key]
ks)) Completions
cs

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

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

-- | Build an '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 forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: TrieLookup -> TrieLookup -> Bool
$c/= :: TrieLookup -> TrieLookup -> Bool
== :: TrieLookup -> TrieLookup -> Bool
$c== :: TrieLookup -> TrieLookup -> Bool
Eq, Eq 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
min :: TrieLookup -> TrieLookup -> TrieLookup
$cmin :: TrieLookup -> TrieLookup -> TrieLookup
max :: TrieLookup -> TrieLookup -> TrieLookup
$cmax :: TrieLookup -> TrieLookup -> TrieLookup
>= :: TrieLookup -> TrieLookup -> Bool
$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
compare :: TrieLookup -> TrieLookup -> Ordering
$ccompare :: TrieLookup -> TrieLookup -> Ordering
Ord, Int -> TrieLookup -> ShowS
[TrieLookup] -> ShowS
TrieLookup -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [TrieLookup] -> ShowS
$cshowList :: [TrieLookup] -> ShowS
show :: TrieLookup -> String
$cshow :: TrieLookup -> String
showsPrec :: Int -> TrieLookup -> ShowS
$cshowsPrec :: Int -> 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 []     = 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
        | forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Key]
ks   -> forall a. a -> Maybe a
Just (Word32 -> TrieLookup
entry Word32
entryOff)
        | Bool
otherwise -> 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       -> forall a. Maybe a
Nothing
        Just Word32
entryOff -> Word32 -> [Key] -> Maybe TrieLookup
go (UArray Word32 Word32
arr 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 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 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (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 forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
! Word32
nodeOff
        keysStart :: Word32
keysStart = Word32
nodeOff forall a. Num a => a -> a -> a
+ Word32
1
        keysEnd :: Word32
keysEnd   = Word32
nodeOff 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 forall a. Ord a => a -> a -> Bool
> Word32
b     = forall a. Maybe a
Nothing
      | Bool
otherwise = case forall a. Ord a => a -> a -> Ordering
compare Word32
key (UArray Word32 Word32
arr 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
midforall a. Num a => a -> a -> a
-Word32
1) Word32
key
          Ordering
EQ -> forall a. a -> Maybe a
Just Word32
mid
          Ordering
GT -> Word32 -> Word32 -> Word32 -> Maybe Word32
bsearch (Word32
midforall a. Num a => a -> a -> a
+Word32
1) Word32
b Word32
key
      where mid :: Word32
mid = (Word32
a forall a. Num a => a -> a -> a
+ Word32
b) 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
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [IntTrieBuilder] -> ShowS
$cshowList :: [IntTrieBuilder] -> ShowS
show :: IntTrieBuilder -> String
$cshow :: IntTrieBuilder -> String
showsPrec :: Int -> IntTrieBuilder -> ShowS
$cshowsPrec :: Int -> IntTrieBuilder -> ShowS
Show, IntTrieBuilder -> IntTrieBuilder -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: IntTrieBuilder -> IntTrieBuilder -> Bool
$c/= :: IntTrieBuilder -> IntTrieBuilder -> Bool
== :: IntTrieBuilder -> IntTrieBuilder -> Bool
$c== :: IntTrieBuilder -> IntTrieBuilder -> Bool
Eq)

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

empty :: IntTrieBuilder
empty :: IntTrieBuilder
empty = IntMap TrieNode -> IntTrieBuilder
IntTrieBuilder 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
  (forall a b. (Integral a, Num b) => a -> b
fromIntegral (Key -> Word32
unKey Key
k) :: Int)
  (forall a b. (a -> b) -> [a] -> [b]
map (forall a b. (Integral a, Num b) => a -> b
fromIntegral 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 forall a b. (a -> b) -> a -> b
$
      forall a. (Maybe a -> Maybe a) -> Int -> IntMap a -> IntMap a
IntMap.alter (\Maybe TrieNode
t' -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$! 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 (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 (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 = 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 forall a b. (a -> b) -> a -> b
$
      forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [e] -> a i e
A.listArray (Word32
0, forall a b. (Integral a, Num b) => a -> b
fromIntegral (IntTrieBuilder -> Int
flatTrieLength IntTrieBuilder
trie) 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 forall a b. (a -> b) -> a -> b
$
        forall a. [(Int, a)] -> IntMap a
IntMap.fromList
          [ (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
  forall a. Num a => a -> a -> a
+ Int
2 forall a. Num a => a -> a -> a
* forall a. IntMap a -> Int
IntMap.size IntMap TrieNode
tns
  forall a. Num a => a -> a -> a
+ forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [ IntTrieBuilder -> Int
flatTrieLength IntTrieBuilder
n | TrieNode IntTrieBuilder
n <- 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 (forall a. [a] -> Q a
queue [IntTrieBuilder
trie]) (IntTrieBuilder -> Int
size IntTrieBuilder
trie)
  where
    size :: IntTrieBuilder -> Int
size (IntTrieBuilder IntMap TrieNode
tns) = Int
1 forall a. Num a => a -> a -> a
+ Int
2 forall a. Num a => a -> a -> a
* 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 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 forall a. [a] -> [a] -> [a]
++ Q IntTrieBuilder -> Int -> [Word32]
go Q IntTrieBuilder
tries' Int
offset'
          where
            !count :: Int
count = forall a. IntMap a -> Int
IntMap.size IntMap TrieNode
tnodes
            flat :: [Word32]
flat   = forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
count
                   forall a. a -> [a] -> [a]
: forall k a. Map k a -> [k]
Map.keys  Map Word32 Word32
keysValues
                  forall a. [a] -> [a] -> [a]
++ forall k a. Map k a -> [a]
Map.elems Map Word32 Word32
keysValues
            (!Int
offset', !Map Word32 Word32
keysValues, !Q IntTrieBuilder
tries') =
              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, 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' = 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 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'   = 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' = 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 = 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 = forall a. [a] -> [a] -> Q a
Q [a]
front (a
x 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)    = forall a. a -> Maybe a
Just (a
x, forall a. [a] -> [a] -> Q a
Q [a]
xs [a]
back)
dequeue (Q []     [a]
back)    = case forall a. [a] -> [a]
reverse [a]
back of
                               a
x:[a]
xs -> forall a. a -> Maybe a
Just (a
x, forall a. [a] -> [a] -> Q a
Q [a]
xs [])
                               []   -> forall a. Maybe a
Nothing

int2Word32 :: Int -> Word32
int2Word32 :: Int -> Word32
int2Word32 = 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) = 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
ixEndforall a. Num a => a -> a -> a
+Word32
1)
 forall a. Semigroup a => a -> a -> a
<> forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Word32
n Builder
r -> Word32 -> Builder
BS.word32BE Word32
n forall a. Semigroup a => a -> a -> a
<> Builder
r) forall a. Monoid a => a
mempty (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) = forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
A.bounds UArray Word32 Word32
arr in
    Int
4
  forall a. Num a => a -> a -> a
+ Int
4 forall a. Num a => a -> a -> a
* (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
ixEnd 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 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 forall a. Num a => a -> a -> a
+ Int
4 forall a. Num a => a -> a -> a
* forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
lenArr
  , ByteString -> Int
BS.length ByteString
bs forall a. Ord a => a -> a -> Bool
>= Int
4 forall a. Num a => a -> a -> a
+ Int
4 forall a. Num a => a -> a -> a
* forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
lenArr
  , let !arr :: UArray Word32 Word32
arr = forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [(i, e)] -> a i e
A.array (Word32
0, Word32
lenArrforall a. Num a => a -> a -> a
-Word32
1)
                      [ (Word32
i, ByteString -> Int -> Word32
readWord32BE ByteString
bs Int
off)
                      | (Word32
i, Int
off) <- forall a b. [a] -> [b] -> [(a, b)]
zip [Word32
0..Word32
lenArrforall a. Num a => a -> a -> a
-Word32
1] [Int
4,Int
8 .. Int
lenTotal forall a. Num a => a -> a -> a
- Int
4] ]
        !bs' :: ByteString
bs' = Int -> ByteString -> ByteString
BS.drop Int
lenTotal ByteString
bs
  = forall a. a -> Maybe a
Just (UArray Word32 Word32 -> IntTrie
IntTrie UArray Word32 Word32
arr, ByteString
bs')

  | Bool
otherwise
  = forall a. Maybe a
Nothing

readWord32BE :: BS.ByteString -> Int -> Word32
readWord32BE :: ByteString -> Int -> Word32
readWord32BE ByteString
bs Int
i =
    forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
i forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
iforall a. Num a => a -> a -> a
+Int
3 forall a. Ord a => a -> a -> Bool
<= ByteString -> Int
BS.length ByteString
bs forall a. Num a => a -> a -> a
- Int
1) forall a b. (a -> b) -> a -> b
$
    forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int -> Word8
BS.unsafeIndex ByteString
bs (Int
i forall a. Num a => a -> a -> a
+ Int
0)) forall a. Bits a => a -> Int -> a
`shiftL` Int
24
  forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int -> Word8
BS.unsafeIndex ByteString
bs (Int
i forall a. Num a => a -> a -> a
+ Int
1)) forall a. Bits a => a -> Int -> a
`shiftL` Int
16
  forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int -> Word8
BS.unsafeIndex ByteString
bs (Int
i forall a. Num a => a -> a -> a
+ Int
2)) forall a. Bits a => a -> Int -> a
`shiftL` Int
8
  forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int -> Word8
BS.unsafeIndex ByteString
bs (Int
i forall a. Num a => a -> a -> a
+ Int
3))