{-# LANGUAGE CPP, BangPatterns, PatternGuards #-}
{-# LANGUAGE DeriveDataTypeable, ScopedTypeVariables #-}
{-# 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)

-- | 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 !arr :: UArray Word32 Word32
arr = (Word32, Word32) -> [(Word32, Word32)] -> UArray Word32 Word32
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [(i, e)] -> a i e
A.array (Word32
0, Word32
lenArrWord32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
-Word32
1)
                      [ (Word32
i, ByteString -> Int -> Word32
readWord32BE ByteString
bs Int
off)
                      | (Word32
i, Int
off) <- [Word32] -> [Int] -> [(Word32, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Word32
0..Word32
lenArrWord32 -> Word32 -> Word32
forall a. Num a => a -> a -> a
-Word32
1] [Int
4,Int
8 .. Int
lenTotal Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
4] ]
        !bs' :: ByteString
bs' = Int -> ByteString -> ByteString
BS.drop Int
lenTotal ByteString
bs
  = (IntTrie, ByteString) -> Maybe (IntTrie, ByteString)
forall a. a -> Maybe a
Just (UArray Word32 Word32 -> IntTrie
IntTrie UArray Word32 Word32
arr, ByteString
bs')

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

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