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
[Trie a] -> ShowS
Trie a -> String
(Int -> Trie a -> ShowS)
-> (Trie a -> String) -> ([Trie a] -> ShowS) -> Show (Trie a)
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
$cshowsPrec :: forall a. Show a => Int -> Trie a -> ShowS
showsPrec :: Int -> Trie a -> ShowS
$cshow :: forall a. Show a => Trie a -> String
show :: Trie a -> String
$cshowList :: forall a. Show a => [Trie a] -> ShowS
showList :: [Trie a] -> ShowS
Show

type Rule = Maybe Int

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

updRule :: Int -> Maybe Int -> Maybe Int
updRule :: Int -> Rule -> Rule
updRule Int
rule = Int -> Rule
forall a. a -> Maybe a
Just (Int -> Rule) -> (Rule -> Int) -> Rule -> Rule
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> (Int -> Int) -> Rule -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
rule (Int -> Int -> Int
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) =
    Rule -> Map Word (Trie Rule) -> Trie Rule
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) =
    Rule -> Map Word (Trie Rule) -> Trie Rule
forall a. a -> Map Word (Trie a) -> Trie a
Branch Rule
rule' ((Maybe (Trie Rule) -> Maybe (Trie Rule))
-> Word -> Map Word (Trie Rule) -> Map Word (Trie Rule)
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
M.alter (Trie Rule -> Maybe (Trie Rule)
forall a. a -> Maybe a
Just (Trie Rule -> Maybe (Trie Rule))
-> (Maybe (Trie Rule) -> Trie Rule)
-> Maybe (Trie Rule)
-> Maybe (Trie Rule)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Trie Rule
-> (Trie Rule -> Trie Rule) -> Maybe (Trie Rule) -> Trie Rule
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 = (Trie Rule -> (Int, String) -> Trie Rule)
-> Trie Rule -> [(Int, String)] -> Trie Rule
forall b a. (b -> a -> b) -> b -> [a] -> b
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 (Char -> [Word]) -> String -> [Word]
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 Map Word (Trie Rule) -> Bool
forall k a. Map k a -> Bool
M.null Map Word (Trie Rule)
ts then
    (Rule, Int) -> Map Word (Trie (Rule, Int)) -> Trie (Rule, Int)
forall a. a -> Map Word (Trie a) -> Trie a
Branch (Rule
rule, Int
0) Map Word (Trie (Rule, Int))
forall a. Monoid a => a
mempty
  else
    let !ts' :: Map Word (Trie (Rule, Int))
ts' = (Trie Rule -> Trie (Rule, Int))
-> Map Word (Trie Rule) -> Map Word (Trie (Rule, Int))
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
    (Rule, Int) -> Map Word (Trie (Rule, Int)) -> Trie (Rule, Int)
forall a. a -> Map Word (Trie a) -> Trie a
Branch (
      Rule
rule,
      Map Word Int -> Int
forall a. Ord a => Map Word a -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum ((Trie (Rule, Int) -> Int)
-> Map Word (Trie (Rule, Int)) -> Map Word Int
forall a b k. (a -> b) -> Map k a -> Map k b
M.map (\(Branch (Rule
rule,Int
d) Map Word (Trie (Rule, Int))
_) -> Int -> (Int -> Int) -> Rule -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Int
d Int -> Int -> Int
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
[Trie' a] -> ShowS
Trie' a -> String
(Int -> Trie' a -> ShowS)
-> (Trie' a -> String) -> ([Trie' a] -> ShowS) -> Show (Trie' a)
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
$cshowsPrec :: forall a. Show a => Int -> Trie' a -> ShowS
showsPrec :: Int -> Trie' a -> ShowS
$cshow :: forall a. Show a => Trie' a -> String
show :: Trie' a -> String
$cshowList :: forall a. Show a => [Trie' a] -> ShowS
showList :: [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 Map Word (Trie (Rule, Int)) -> [(Word, Trie (Rule, Int))]
forall k a. Map k a -> [(k, a)]
M.toList Map Word (Trie (Rule, Int))
ts of
  [] -> (Rule, Int) -> Map Word (Trie' (Rule, Int)) -> Trie' (Rule, Int)
forall a. a -> Map Word (Trie' a) -> Trie' a
Branch' (Rule, Int)
a Map Word (Trie' (Rule, Int))
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 -> (Rule, Int) -> [Word] -> Trie' (Rule, Int) -> Trie' (Rule, Int)
forall a. a -> [Word] -> Trie' a -> Trie' a
Path (Rule, Int)
a (Word
wWord -> [Word] -> [Word]
forall a. a -> [a] -> [a]
:[Word]
ws) Trie' (Rule, Int)
t
           Trie' (Rule, Int)
t                      -> (Rule, Int) -> [Word] -> Trie' (Rule, Int) -> Trie' (Rule, Int)
forall a. a -> [Word] -> Trie' a -> Trie' a
Path (Rule, Int)
a [Word
w] Trie' (Rule, Int)
t
  [(Word, Trie (Rule, Int))]
_   -> (Rule, Int) -> Map Word (Trie' (Rule, Int)) -> Trie' (Rule, Int)
forall a. a -> Map Word (Trie' a) -> Trie' a
Branch' (Rule, Int)
a ((Trie (Rule, Int) -> Trie' (Rule, Int))
-> Map Word (Trie (Rule, Int)) -> Map Word (Trie' (Rule, Int))
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 Rule
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)
    | Map Word (Trie' (Rule, Int)) -> Bool
forall k a. Map k a -> Bool
M.null Map Word (Trie' (Rule, Int))
ts        = (Rule, Int, Int)
-> Map Word (Trie' (Rule, Int, Int)) -> Trie' (Rule, Int, Int)
forall a. a -> Map Word (Trie' a) -> Trie' a
Branch' (Rule
rule', Int
0, Int
d) Map Word (Trie' (Rule, Int, Int))
forall a. Monoid a => a
mempty
    | Rule
Nothing <- Rule
rule' = (Rule, Int, Int)
-> Map Word (Trie' (Rule, Int, Int)) -> Trie' (Rule, Int, Int)
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Trie' (Rule, Int) -> Trie' (Rule, Int, Int))
-> Map Word (Trie' (Rule, Int))
-> Map Word (Trie' (Rule, Int, Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Word (Trie' (Rule, Int))
ts)
    | Bool
otherwise        = (Rule, Int, Int)
-> Map Word (Trie' (Rule, Int, Int)) -> Trie' (Rule, Int, Int)
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     (Trie' (Rule, Int) -> Trie' (Rule, Int, Int))
-> Map Word (Trie' (Rule, Int))
-> Map Word (Trie' (Rule, Int, Int))
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' = (Rule, Int, Int)
-> [Word] -> Trie' (Rule, Int, Int) -> Trie' (Rule, Int, Int)
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Word] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Word]
ws) Trie' (Rule, Int)
t)
    | Bool
otherwise        = (Rule, Int, Int)
-> [Word] -> Trie' (Rule, Int, Int) -> Trie' (Rule, Int, Int)
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' ([Word] -> Int
forall a. [a] -> Int
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
      | Map Word (Trie' (Rule, Int, Int)) -> Bool
forall k a. Map k a -> Bool
M.null Map Word (Trie' (Rule, Int, Int))
ts -> (Rule, Int, Rule)
-> Map Word (Trie' (Rule, Int, Rule)) -> Trie' (Rule, Int, Rule)
forall a. a -> Map Word (Trie' a) -> Trie' a
Branch' (Rule
r, Int
n, Rule
forall a. Maybe a
Nothing) Map Word (Trie' (Rule, Int, Rule))
forall a. Monoid a => a
mempty
      | Int
res Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
1   -> (Rule, Int, Rule)
-> Map Word (Trie' (Rule, Int, Rule)) -> Trie' (Rule, Int, Rule)
forall a. a -> Map Word (Trie' a) -> Trie' a
Branch' (Rule
r, Int
n, Int -> Rule
forall a. a -> Maybe a
Just Int
d ) (Int -> Trie' (Rule, Int, Int) -> Trie' (Rule, Int, Rule)
go (Int
d   Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Trie' (Rule, Int, Int) -> Trie' (Rule, Int, Rule))
-> Map Word (Trie' (Rule, Int, Int))
-> Map Word (Trie' (Rule, Int, Rule))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map Word (Trie' (Rule, Int, Int))
ts)
      | Bool
otherwise -> (Rule, Int, Rule)
-> Map Word (Trie' (Rule, Int, Rule)) -> Trie' (Rule, Int, Rule)
forall a. a -> Map Word (Trie' a) -> Trie' a
Branch' (Rule
r, Int
n, Rule
forall a. Maybe a
Nothing) (Int -> Trie' (Rule, Int, Int) -> Trie' (Rule, Int, Rule)
go (Int
res Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Trie' (Rule, Int, Int) -> Trie' (Rule, Int, Rule))
-> Map Word (Trie' (Rule, Int, Int))
-> Map Word (Trie' (Rule, Int, Rule))
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 [Word] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Word]
ws of
      Int
l | Int
res Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
l   -> (Rule, Int, Rule)
-> [Word] -> Trie' (Rule, Int, Rule) -> Trie' (Rule, Int, Rule)
forall a. a -> [Word] -> Trie' a -> Trie' a
Path (Rule
r, Int
n, Int -> Rule
forall a. a -> Maybe a
Just (Int -> Rule) -> Int -> Rule
forall a b. (a -> b) -> a -> b
$! Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
res) [Word]
ws (Int -> Trie' (Rule, Int, Int) -> Trie' (Rule, Int, Rule)
go (Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l)   Trie' (Rule, Int, Int)
t)
        | Bool
otherwise -> (Rule, Int, Rule)
-> [Word] -> Trie' (Rule, Int, Rule) -> Trie' (Rule, Int, Rule)
forall a. a -> [Word] -> Trie' a -> Trie' a
Path (Rule
r, Int
n, Rule
forall a. Maybe a
Nothing        ) [Word]
ws (Int -> Trie' (Rule, Int, Int) -> Trie' (Rule, Int, Rule)
go (Int
res Int -> Int -> Int
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 (Trie' (Rule, Int, Int) -> Trie' (Rule, Int, Rule))
-> ([(Int, String)] -> Trie' (Rule, Int, Int))
-> [(Int, String)]
-> Trie' (Rule, Int, Rule)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Trie' (Rule, Int) -> Trie' (Rule, Int, Int)
fallbacks (Trie' (Rule, Int) -> Trie' (Rule, Int, Int))
-> ([(Int, String)] -> Trie' (Rule, Int))
-> [(Int, String)]
-> Trie' (Rule, Int, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Trie (Rule, Int) -> Trie' (Rule, Int)
pathify (Trie (Rule, Int) -> Trie' (Rule, Int))
-> ([(Int, String)] -> Trie (Rule, Int))
-> [(Int, String)]
-> Trie' (Rule, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Trie Rule -> Trie (Rule, Int)
mindepths (Trie Rule -> Trie (Rule, Int))
-> ([(Int, String)] -> Trie Rule)
-> [(Int, String)]
-> Trie (Rule, Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Int, String)] -> Trie Rule
listToTrie