-- | Efficient literal branching using Template Haskell.

module FlatParse.Basic.Switch
  ( switch, switchWithPost, rawSwitchWithPost
  ) where

import Control.Monad
import Data.Foldable
import Data.Map (Map)
import Language.Haskell.TH

import qualified Data.Map.Strict as M

import FlatParse.Common.Switch
import FlatParse.Basic.Base ( ensure, skipBack, branch, failed )
import FlatParse.Basic.Bytes ( bytesUnsafe )
import FlatParse.Basic.Integers ( anyWord8Unsafe )

{-|
This is a template function which makes it possible to branch on a collection of string literals in
an efficient way. By using `switch`, such branching is compiled to a trie of primitive parsing
operations, which has optimized control flow, vectorized reads and grouped checking for needed input
bytes.

The syntax is slightly magical, it overloads the usual @case@ expression. An example:

@
    $(switch [| case _ of
        "foo" -> pure True
        "bar" -> pure False |])
@

The underscore is mandatory in @case _ of@. Each branch must be a string literal, but optionally
we may have a default case, like in

@
    $(switch [| case _ of
        "foo" -> pure 10
        "bar" -> pure 20
        _     -> pure 30 |])
@

All case right hand sides must be parsers with the same type. That type is also the type
of the whole `switch` expression.

A `switch` has longest match semantics, and the order of cases does not matter, except for
the default case, which may only appear as the last case.

If a `switch` does not have a default case, and no case matches the input, then it returns with
failure, \without\ having consumed any input. A fallthrough to the default case also does not
consume any input.
-}
switch :: Q Exp -> Q Exp
switch :: Q Exp -> Q Exp
switch = Maybe (Q Exp) -> Q Exp -> Q Exp
switchWithPost Maybe (Q Exp)
forall a. Maybe a
Nothing

{-|
Switch expression with an optional first argument for performing a post-processing action after
every successful branch matching, not including the default branch. For example, if we have
@ws :: ParserT st e ()@ for a whitespace parser, we might want to consume whitespace after matching
on any of the switch cases. For that case, we can define a "lexeme" version of `switch` as
follows.

@
  switch' :: Q Exp -> Q Exp
  switch' = switchWithPost (Just [| ws |])
@

Note that this @switch'@ function cannot be used in the same module it's defined in, because of the
stage restriction of Template Haskell.
-}
switchWithPost :: Maybe (Q Exp) -> Q Exp -> Q Exp
switchWithPost :: Maybe (Q Exp) -> Q Exp -> Q Exp
switchWithPost Maybe (Q Exp)
postAction Q Exp
exp = do
  !Maybe Exp
postAction <- Maybe (Q Exp) -> Q (Maybe Exp)
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => Maybe (m a) -> m (Maybe a)
sequence Maybe (Q Exp)
postAction
  (![(String, Exp)]
cases, !Maybe Exp
fallback) <- Q Exp -> Q ([(String, Exp)], Maybe Exp)
parseSwitch Q Exp
exp
  (Map (Maybe Int) Exp, Trie' (Maybe Int, Int, Maybe Int)) -> Q Exp
genTrie ((Map (Maybe Int) Exp, Trie' (Maybe Int, Int, Maybe Int)) -> Q Exp)
-> (Map (Maybe Int) Exp, Trie' (Maybe Int, Int, Maybe Int))
-> Q Exp
forall a b. (a -> b) -> a -> b
$! Maybe Exp
-> [(String, Exp)]
-> Maybe Exp
-> (Map (Maybe Int) Exp, Trie' (Maybe Int, Int, Maybe Int))
genSwitchTrie' Maybe Exp
postAction [(String, Exp)]
cases Maybe Exp
fallback

-- | Version of `switchWithPost` without syntactic sugar. The second argument is the
--   list of cases, the third is the default case.
rawSwitchWithPost :: Maybe (Q Exp) -> [(String, Q Exp)] -> Maybe (Q Exp) -> Q Exp
rawSwitchWithPost :: Maybe (Q Exp) -> [(String, Q Exp)] -> Maybe (Q Exp) -> Q Exp
rawSwitchWithPost Maybe (Q Exp)
postAction [(String, Q Exp)]
cases Maybe (Q Exp)
fallback = do
  !Maybe Exp
postAction <- Maybe (Q Exp) -> Q (Maybe Exp)
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => Maybe (m a) -> m (Maybe a)
sequence Maybe (Q Exp)
postAction
  ![(String, Exp)]
cases <- [(String, Q Exp)]
-> ((String, Q Exp) -> Q (String, Exp)) -> Q [(String, Exp)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(String, Q Exp)]
cases \(String
str, Q Exp
rhs) -> (String
str,) (Exp -> (String, Exp)) -> Q Exp -> Q (String, Exp)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Q Exp
rhs
  !Maybe Exp
fallback <- Maybe (Q Exp) -> Q (Maybe Exp)
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => Maybe (m a) -> m (Maybe a)
sequence Maybe (Q Exp)
fallback
  (Map (Maybe Int) Exp, Trie' (Maybe Int, Int, Maybe Int)) -> Q Exp
genTrie ((Map (Maybe Int) Exp, Trie' (Maybe Int, Int, Maybe Int)) -> Q Exp)
-> (Map (Maybe Int) Exp, Trie' (Maybe Int, Int, Maybe Int))
-> Q Exp
forall a b. (a -> b) -> a -> b
$! Maybe Exp
-> [(String, Exp)]
-> Maybe Exp
-> (Map (Maybe Int) Exp, Trie' (Maybe Int, Int, Maybe Int))
genSwitchTrie' Maybe Exp
postAction [(String, Exp)]
cases Maybe Exp
fallback

#if MIN_VERSION_base(4,15,0)
mkDoE :: [Stmt] -> Exp
mkDoE = Maybe ModName -> [Stmt] -> Exp
DoE Maybe ModName
forall a. Maybe a
Nothing
{-# inline mkDoE #-}
#else
mkDoE = DoE
{-# inline mkDoE #-}
#endif

genTrie :: (Map (Maybe Int) Exp, Trie' (Rule, Int, Maybe Int)) -> Q Exp
genTrie :: (Map (Maybe Int) Exp, Trie' (Maybe Int, Int, Maybe Int)) -> Q Exp
genTrie (Map (Maybe Int) Exp
rules, Trie' (Maybe Int, Int, Maybe Int)
t) = do
  Map (Maybe Int) (Name, Exp)
branches <- (Exp -> Q (Name, Exp))
-> Map (Maybe Int) Exp -> Q (Map (Maybe Int) (Name, Exp))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map (Maybe Int) a -> f (Map (Maybe Int) b)
traverse (\Exp
e -> (,) (Name -> Exp -> (Name, Exp)) -> Q Name -> Q (Exp -> (Name, Exp))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"rule") Q (Exp -> (Name, Exp)) -> Q Exp -> Q (Name, Exp)
forall a b. Q (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Exp -> Q Exp
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
e) Map (Maybe Int) Exp
rules

  let ix :: Map a a -> a -> a
ix Map a a
m a
k = case a -> Map a a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
k Map a a
m of
        Maybe a
Nothing -> String -> a
forall a. HasCallStack => String -> a
error (String
"key not in map: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Show a => a -> String
show a
k)
        Just a
a  -> a
a

  let ensure' :: Maybe Int -> Maybe (Q Exp)
      ensure' :: Maybe Int -> Maybe (Q Exp)
ensure' = (Int -> Q Exp) -> Maybe Int -> Maybe (Q Exp)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Int
n -> [| ensure n |])

      fallback :: Rule -> Int ->  Q Exp
      fallback :: Maybe Int -> Int -> Q Exp
fallback Maybe Int
rule Int
0 = Exp -> Q Exp
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> Q Exp) -> Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ Name -> Exp
VarE (Name -> Exp) -> Name -> Exp
forall a b. (a -> b) -> a -> b
$ (Name, Exp) -> Name
forall a b. (a, b) -> a
fst ((Name, Exp) -> Name) -> (Name, Exp) -> Name
forall a b. (a -> b) -> a -> b
$ Map (Maybe Int) (Name, Exp) -> Maybe Int -> (Name, Exp)
forall {a} {a}. (Ord a, Show a) => Map a a -> a -> a
ix Map (Maybe Int) (Name, Exp)
branches Maybe Int
rule
      fallback Maybe Int
rule Int
n = [| skipBack n >> $(Exp -> Q Exp
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> Q Exp) -> Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ Name -> Exp
VarE (Name -> Exp) -> Name -> Exp
forall a b. (a -> b) -> a -> b
$ (Name, Exp) -> Name
forall a b. (a, b) -> a
fst ((Name, Exp) -> Name) -> (Name, Exp) -> Name
forall a b. (a -> b) -> a -> b
$ Map (Maybe Int) (Name, Exp) -> Maybe Int -> (Name, Exp)
forall {a} {a}. (Ord a, Show a) => Map a a -> a -> a
ix Map (Maybe Int) (Name, Exp)
branches Maybe Int
rule) |]

  let go :: Trie' (Rule, Int, Maybe Int) -> Q Exp
      go :: Trie' (Maybe Int, Int, Maybe Int) -> Q Exp
go = \case
        Branch' (Maybe Int
r, Int
n, Maybe Int
alloc) Map Word (Trie' (Maybe Int, Int, Maybe Int))
ts
          | Map Word (Trie' (Maybe Int, Int, Maybe Int)) -> Bool
forall k a. Map k a -> Bool
M.null Map Word (Trie' (Maybe Int, Int, Maybe Int))
ts -> Exp -> Q Exp
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> Q Exp) -> Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ Name -> Exp
VarE (Name -> Exp) -> Name -> Exp
forall a b. (a -> b) -> a -> b
$ (Name, Exp) -> Name
forall a b. (a, b) -> a
fst ((Name, Exp) -> Name) -> (Name, Exp) -> Name
forall a b. (a -> b) -> a -> b
$ Map (Maybe Int) (Name, Exp)
branches Map (Maybe Int) (Name, Exp) -> Maybe Int -> (Name, Exp)
forall k a. Ord k => Map k a -> k -> a
M.! Maybe Int
r
          | Bool
otherwise -> do
              ![(Word, Exp)]
next         <- (((Word, Trie' (Maybe Int, Int, Maybe Int)) -> Q (Word, Exp))
-> [(Word, Trie' (Maybe Int, Int, Maybe Int))] -> Q [(Word, Exp)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (((Word, Trie' (Maybe Int, Int, Maybe Int)) -> Q (Word, Exp))
 -> [(Word, Trie' (Maybe Int, Int, Maybe Int))] -> Q [(Word, Exp)])
-> ((Trie' (Maybe Int, Int, Maybe Int) -> Q Exp)
    -> (Word, Trie' (Maybe Int, Int, Maybe Int)) -> Q (Word, Exp))
-> (Trie' (Maybe Int, Int, Maybe Int) -> Q Exp)
-> [(Word, Trie' (Maybe Int, Int, Maybe Int))]
-> Q [(Word, Exp)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Trie' (Maybe Int, Int, Maybe Int) -> Q Exp)
-> (Word, Trie' (Maybe Int, Int, Maybe Int)) -> Q (Word, Exp)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> (Word, a) -> f (Word, b)
traverse) Trie' (Maybe Int, Int, Maybe Int) -> Q Exp
go (Map Word (Trie' (Maybe Int, Int, Maybe Int))
-> [(Word, Trie' (Maybe Int, Int, Maybe Int))]
forall k a. Map k a -> [(k, a)]
M.toList Map Word (Trie' (Maybe Int, Int, Maybe Int))
ts)
              !Exp
defaultCase  <- Maybe Int -> Int -> Q Exp
fallback Maybe Int
r (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

              let cases :: Exp
cases = [Stmt] -> Exp
mkDoE ([Stmt] -> Exp) -> [Stmt] -> Exp
forall a b. (a -> b) -> a -> b
$
                    [Pat -> Exp -> Stmt
BindS (Name -> Pat
VarP (String -> Name
mkName String
"c")) (Name -> Exp
VarE 'anyWord8Unsafe),
                      Exp -> Stmt
NoBindS (Exp -> [Match] -> Exp
CaseE (Name -> Exp
VarE (String -> Name
mkName String
"c"))
                         (((Word, Exp) -> Match) -> [(Word, Exp)] -> [Match]
forall a b. (a -> b) -> [a] -> [b]
map (\(Word
w, Exp
t) ->
                                 Pat -> Body -> [Dec] -> Match
Match (Lit -> Pat
LitP (Integer -> Lit
IntegerL (Word -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
w)))
                                       (Exp -> Body
NormalB Exp
t)
                                       [])
                              [(Word, Exp)]
next
                          [Match] -> [Match] -> [Match]
forall a. [a] -> [a] -> [a]
++ [Pat -> Body -> [Dec] -> Match
Match Pat
WildP (Exp -> Body
NormalB Exp
defaultCase) []]))]

              case Maybe Int -> Maybe (Q Exp)
ensure' Maybe Int
alloc of
                Maybe (Q Exp)
Nothing    -> Exp -> Q Exp
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
cases
                Just Q Exp
alloc -> [| branch $Q Exp
alloc $(Exp -> Q Exp
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
cases) $(Maybe Int -> Int -> Q Exp
fallback Maybe Int
r Int
n) |]

        Path (Maybe Int
r, Int
n, Maybe Int
alloc) [Word]
ws Trie' (Maybe Int, Int, Maybe Int)
t ->
          case Maybe Int -> Maybe (Q Exp)
ensure' Maybe Int
alloc of
            Maybe (Q Exp)
Nothing    -> [| branch $([Word] -> Q Exp
bytesUnsafe [Word]
ws) $(Trie' (Maybe Int, Int, Maybe Int) -> Q Exp
go Trie' (Maybe Int, Int, Maybe Int)
t) $(Maybe Int -> Int -> Q Exp
fallback Maybe Int
r Int
n)|]
            Just Q Exp
alloc -> [| branch ($Q Exp
alloc >> $([Word] -> Q Exp
bytesUnsafe [Word]
ws)) $(Trie' (Maybe Int, Int, Maybe Int) -> Q Exp
go Trie' (Maybe Int, Int, Maybe Int)
t) $(Maybe Int -> Int -> Q Exp
fallback Maybe Int
r Int
n) |]

  [Q Dec] -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => [m Dec] -> m Exp -> m Exp
letE
    (((Name, Exp) -> Q Dec) -> [(Name, Exp)] -> [Q Dec]
forall a b. (a -> b) -> [a] -> [b]
map (\(Name
x, Exp
rhs) -> Q Pat -> Q Body -> [Q Dec] -> Q Dec
forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
x) (Q Exp -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB (Exp -> Q Exp
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
rhs)) []) (Map (Maybe Int) (Name, Exp) -> [(Name, Exp)]
forall a. Map (Maybe Int) a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
Data.Foldable.toList Map (Maybe Int) (Name, Exp)
branches))
    (Trie' (Maybe Int, Int, Maybe Int) -> Q Exp
go Trie' (Maybe Int, Int, Maybe Int)
t)

genSwitchTrie' :: Maybe Exp -> [(String, Exp)] -> Maybe Exp
              -> (Map (Maybe Int) Exp, Trie' (Rule, Int, Maybe Int))
genSwitchTrie' :: Maybe Exp
-> [(String, Exp)]
-> Maybe Exp
-> (Map (Maybe Int) Exp, Trie' (Maybe Int, Int, Maybe Int))
genSwitchTrie' Maybe Exp
postAction [(String, Exp)]
cases Maybe Exp
fallback =

  let (![(Maybe Int, Exp)]
branches, ![(Int, String)]
strings) = [((Maybe Int, Exp), (Int, String))]
-> ([(Maybe Int, Exp)], [(Int, String)])
forall a b. [(a, b)] -> ([a], [b])
unzip do
        (!Int
i, (!String
str, !Exp
rhs)) <- [Int] -> [(String, Exp)] -> [(Int, (String, Exp))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] [(String, Exp)]
cases
        case Maybe Exp
postAction of
          Maybe Exp
Nothing    -> ((Maybe Int, Exp), (Int, String))
-> [((Maybe Int, Exp), (Int, String))]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Int -> Maybe Int
forall a. a -> Maybe a
Just Int
i, Exp
rhs), (Int
i, String
str))
          Just !Exp
post -> ((Maybe Int, Exp), (Int, String))
-> [((Maybe Int, Exp), (Int, String))]
forall a. a -> [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Int -> Maybe Int
forall a. a -> Maybe a
Just Int
i, (Name -> Exp
VarE '(>>)) Exp -> Exp -> Exp
`AppE` Exp
post Exp -> Exp -> Exp
`AppE` Exp
rhs), (Int
i, String
str))

      !m :: Map (Maybe Int) Exp
m    = [(Maybe Int, Exp)] -> Map (Maybe Int) Exp
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ((Maybe Int
forall a. Maybe a
Nothing, Exp -> (Exp -> Exp) -> Maybe Exp -> Exp
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Name -> Exp
VarE 'failed) Exp -> Exp
forall a. a -> a
id Maybe Exp
fallback) (Maybe Int, Exp) -> [(Maybe Int, Exp)] -> [(Maybe Int, Exp)]
forall a. a -> [a] -> [a]
: [(Maybe Int, Exp)]
branches)
      !trie :: Trie' (Maybe Int, Int, Maybe Int)
trie = [(Int, String)] -> Trie' (Maybe Int, Int, Maybe Int)
compileTrie [(Int, String)]
strings
  in (Map (Maybe Int) Exp
m , Trie' (Maybe Int, Int, Maybe Int)
trie)