module FlatParse.Common.Switch where

import Data.Foldable (foldl')
import Data.Map (Map)

import qualified Data.Map.Strict as M

import FlatParse.Common.Assorted

-- Switch trie compilation
--------------------------------------------------------------------------------

data Trie a = Branch !a !(Map Word (Trie a))
  deriving Int -> Trie a -> ShowS
forall a. Show a => Int -> Trie a -> ShowS
forall a. Show a => [Trie a] -> ShowS
forall a. Show a => Trie a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Trie a] -> ShowS
$cshowList :: forall a. Show a => [Trie a] -> ShowS
show :: Trie a -> String
$cshow :: forall a. Show a => Trie a -> String
showsPrec :: Int -> Trie a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Trie a -> ShowS
Show

type Rule = Maybe Int

nilTrie :: Trie Rule
nilTrie :: Trie Rule
nilTrie = forall a. a -> Map Word (Trie a) -> Trie a
Branch forall a. Maybe a
Nothing forall a. Monoid a => a
mempty

updRule :: Int -> Maybe Int -> Maybe Int
updRule :: Int -> Rule -> Rule
updRule Int
rule = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
rule (forall a. Ord a => a -> a -> a
min Int
rule)

insert :: Int -> [Word] -> Trie Rule -> Trie Rule
insert :: Int -> [Word] -> Trie Rule -> Trie Rule
insert Int
rule = [Word] -> Trie Rule -> Trie Rule
go where
  go :: [Word] -> Trie Rule -> Trie Rule
go [] (Branch Rule
rule' Map Word (Trie Rule)
ts) =
    forall a. a -> Map Word (Trie a) -> Trie a
Branch (Int -> Rule -> Rule
updRule Int
rule Rule
rule') Map Word (Trie Rule)
ts
  go (Word
c:[Word]
cs) (Branch Rule
rule' Map Word (Trie Rule)
ts) =
    forall a. a -> Map Word (Trie a) -> Trie a
Branch Rule
rule' (forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
M.alter (forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall b a. b -> (a -> b) -> Maybe a -> b
maybe ([Word] -> Trie Rule -> Trie Rule
go [Word]
cs Trie Rule
nilTrie) ([Word] -> Trie Rule -> Trie Rule
go [Word]
cs)) Word
c Map Word (Trie Rule)
ts)

listToTrie :: [(Int, String)] -> Trie Rule
listToTrie :: [(Int, String)] -> Trie Rule
listToTrie = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Trie Rule
t (!Int
r, !String
s) -> Int -> [Word] -> Trie Rule -> Trie Rule
insert Int
r (Char -> [Word]
charToBytes forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String
s) Trie Rule
t) Trie Rule
nilTrie

-- | Decorate a trie with the minimum lengths of non-empty paths. This
--   is used later to place `ensureBytes#`.
mindepths :: Trie Rule -> Trie (Rule, Int)
mindepths :: Trie Rule -> Trie (Rule, Int)
mindepths (Branch Rule
rule Map Word (Trie Rule)
ts) =
  if forall k a. Map k a -> Bool
M.null Map Word (Trie Rule)
ts then
    forall a. a -> Map Word (Trie a) -> Trie a
Branch (Rule
rule, Int
0) forall a. Monoid a => a
mempty
  else
    let !ts' :: Map Word (Trie (Rule, Int))
ts' = forall a b k. (a -> b) -> Map k a -> Map k b
M.map Trie Rule -> Trie (Rule, Int)
mindepths Map Word (Trie Rule)
ts in
    forall a. a -> Map Word (Trie a) -> Trie a
Branch (
      Rule
rule,
      forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum (forall a b k. (a -> b) -> Map k a -> Map k b
M.map (\(Branch (Rule
rule,Int
d) Map Word (Trie (Rule, Int))
_) -> forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Int
d forall a. Num a => a -> a -> a
+ Int
1) (\Int
_ -> Int
1) Rule
rule) Map Word (Trie (Rule, Int))
ts'))
      Map Word (Trie (Rule, Int))
ts'

data Trie' a
  = Branch' !a !(Map Word (Trie' a))
  | Path !a ![Word] !(Trie' a)
  deriving Int -> Trie' a -> ShowS
forall a. Show a => Int -> Trie' a -> ShowS
forall a. Show a => [Trie' a] -> ShowS
forall a. Show a => Trie' a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Trie' a] -> ShowS
$cshowList :: forall a. Show a => [Trie' a] -> ShowS
show :: Trie' a -> String
$cshow :: forall a. Show a => Trie' a -> String
showsPrec :: Int -> Trie' a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Trie' a -> ShowS
Show

-- | Compress linear paths.
pathify :: Trie (Rule, Int) -> Trie' (Rule, Int)
pathify :: Trie (Rule, Int) -> Trie' (Rule, Int)
pathify (Branch (Rule, Int)
a Map Word (Trie (Rule, Int))
ts) = case forall k a. Map k a -> [(k, a)]
M.toList Map Word (Trie (Rule, Int))
ts of
  [] -> forall a. a -> Map Word (Trie' a) -> Trie' a
Branch' (Rule, Int)
a forall a. Monoid a => a
mempty
  [(Word
w, Trie (Rule, Int)
t)] -> case Trie (Rule, Int) -> Trie' (Rule, Int)
pathify Trie (Rule, Int)
t of
           Path (Rule
Nothing, Int
_) [Word]
ws Trie' (Rule, Int)
t -> forall a. a -> [Word] -> Trie' a -> Trie' a
Path (Rule, Int)
a (Word
wforall a. a -> [a] -> [a]
:[Word]
ws) Trie' (Rule, Int)
t
           Trie' (Rule, Int)
t                      -> forall a. a -> [Word] -> Trie' a -> Trie' a
Path (Rule, Int)
a [Word
w] Trie' (Rule, Int)
t
  [(Word, Trie (Rule, Int))]
_   -> forall a. a -> Map Word (Trie' a) -> Trie' a
Branch' (Rule, Int)
a (forall a b k. (a -> b) -> Map k a -> Map k b
M.map Trie (Rule, Int) -> Trie' (Rule, Int)
pathify Map Word (Trie (Rule, Int))
ts)

-- | Compute where to fall back after we exhausted a branch. If the branch is
--   empty, that means we've succeded at reading and we jump to the rhs rule.
fallbacks :: Trie' (Rule, Int) -> Trie' (Rule, Int, Int)
fallbacks :: Trie' (Rule, Int) -> Trie' (Rule, Int, Int)
fallbacks = Rule -> Int -> Trie' (Rule, Int) -> Trie' (Rule, Int, Int)
go forall a. Maybe a
Nothing Int
0  where
  go :: Rule -> Int -> Trie' (Rule, Int) -> Trie' (Rule, Int, Int)
  go :: Rule -> Int -> Trie' (Rule, Int) -> Trie' (Rule, Int, Int)
go !Rule
rule !Int
n (Branch' (Rule
rule', Int
d) Map Word (Trie' (Rule, Int))
ts)
    | forall k a. Map k a -> Bool
M.null Map Word (Trie' (Rule, Int))
ts        = forall a. a -> Map Word (Trie' a) -> Trie' a
Branch' (Rule
rule', Int
0, Int
d) forall a. Monoid a => a
mempty
    | Rule
Nothing <- Rule
rule' = forall a. a -> Map Word (Trie' a) -> Trie' a
Branch' (Rule
rule, Int
n, Int
d) (Rule -> Int -> Trie' (Rule, Int) -> Trie' (Rule, Int, Int)
go Rule
rule (Int
n forall a. Num a => a -> a -> a
+ Int
1) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Word (Trie' (Rule, Int))
ts)
    | Bool
otherwise        = forall a. a -> Map Word (Trie' a) -> Trie' a
Branch' (Rule
rule', Int
0, Int
d) (Rule -> Int -> Trie' (Rule, Int) -> Trie' (Rule, Int, Int)
go Rule
rule' Int
1     forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Word (Trie' (Rule, Int))
ts)
  go Rule
rule Int
n (Path (Rule
rule', Int
d) [Word]
ws Trie' (Rule, Int)
t)
    | Rule
Nothing <- Rule
rule' = forall a. a -> [Word] -> Trie' a -> Trie' a
Path (Rule
rule, Int
n, Int
d)  [Word]
ws (Rule -> Int -> Trie' (Rule, Int) -> Trie' (Rule, Int, Int)
go Rule
rule (Int
n forall a. Num a => a -> a -> a
+ forall (t :: * -> *) a. Foldable t => t a -> Int
length [Word]
ws) Trie' (Rule, Int)
t)
    | Bool
otherwise        = forall a. a -> [Word] -> Trie' a -> Trie' a
Path (Rule
rule', Int
0, Int
d) [Word]
ws (Rule -> Int -> Trie' (Rule, Int) -> Trie' (Rule, Int, Int)
go Rule
rule' (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Word]
ws) Trie' (Rule, Int)
t)

-- | Decorate with `ensureBytes#` invocations, represented as
--   `Maybe Int`.
ensureBytes :: Trie' (Rule, Int, Int) -> Trie' (Rule, Int, Maybe Int)
ensureBytes :: Trie' (Rule, Int, Int) -> Trie' (Rule, Int, Rule)
ensureBytes = Int -> Trie' (Rule, Int, Int) -> Trie' (Rule, Int, Rule)
go Int
0 where
  go :: Int -> Trie' (Rule, Int, Int) -> Trie' (Rule, Int, Maybe Int)
  go :: Int -> Trie' (Rule, Int, Int) -> Trie' (Rule, Int, Rule)
go !Int
res = \case
    Branch' (Rule
r, Int
n, Int
d) Map Word (Trie' (Rule, Int, Int))
ts
      | forall k a. Map k a -> Bool
M.null Map Word (Trie' (Rule, Int, Int))
ts -> forall a. a -> Map Word (Trie' a) -> Trie' a
Branch' (Rule
r, Int
n, forall a. Maybe a
Nothing) forall a. Monoid a => a
mempty
      | Int
res forall a. Ord a => a -> a -> Bool
< Int
1   -> forall a. a -> Map Word (Trie' a) -> Trie' a
Branch' (Rule
r, Int
n, forall a. a -> Maybe a
Just Int
d ) (Int -> Trie' (Rule, Int, Int) -> Trie' (Rule, Int, Rule)
go (Int
d   forall a. Num a => a -> a -> a
- Int
1) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Word (Trie' (Rule, Int, Int))
ts)
      | Bool
otherwise -> forall a. a -> Map Word (Trie' a) -> Trie' a
Branch' (Rule
r, Int
n, forall a. Maybe a
Nothing) (Int -> Trie' (Rule, Int, Int) -> Trie' (Rule, Int, Rule)
go (Int
res forall a. Num a => a -> a -> a
- Int
1) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Word (Trie' (Rule, Int, Int))
ts)
    Path (Rule
r, Int
n, Int
d) [Word]
ws Trie' (Rule, Int, Int)
t -> case forall (t :: * -> *) a. Foldable t => t a -> Int
length [Word]
ws of
      Int
l | Int
res forall a. Ord a => a -> a -> Bool
< Int
l   -> forall a. a -> [Word] -> Trie' a -> Trie' a
Path (Rule
r, Int
n, forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$! Int
d forall a. Num a => a -> a -> a
- Int
res) [Word]
ws (Int -> Trie' (Rule, Int, Int) -> Trie' (Rule, Int, Rule)
go (Int
d forall a. Num a => a -> a -> a
- Int
l)   Trie' (Rule, Int, Int)
t)
        | Bool
otherwise -> forall a. a -> [Word] -> Trie' a -> Trie' a
Path (Rule
r, Int
n, forall a. Maybe a
Nothing        ) [Word]
ws (Int -> Trie' (Rule, Int, Int) -> Trie' (Rule, Int, Rule)
go (Int
res forall a. Num a => a -> a -> a
- Int
l) Trie' (Rule, Int, Int)
t)

compileTrie :: [(Int, String)] -> Trie' (Rule, Int, Maybe Int)
compileTrie :: [(Int, String)] -> Trie' (Rule, Int, Rule)
compileTrie = Trie' (Rule, Int, Int) -> Trie' (Rule, Int, Rule)
ensureBytes forall b c a. (b -> c) -> (a -> b) -> a -> c
. Trie' (Rule, Int) -> Trie' (Rule, Int, Int)
fallbacks forall b c a. (b -> c) -> (a -> b) -> a -> c
. Trie (Rule, Int) -> Trie' (Rule, Int)
pathify forall b c a. (b -> c) -> (a -> b) -> a -> c
. Trie Rule -> Trie (Rule, Int)
mindepths forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Int, String)] -> Trie Rule
listToTrie