-- | 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 ( 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 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 <- forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t 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 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 <- forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence Maybe (Q Exp)
postAction
  ![(String, Exp)]
cases <- 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,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Q Exp
rhs
  !Maybe Exp
fallback <- forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence Maybe (Q Exp)
fallback
  (Map (Maybe Int) Exp, Trie' (Maybe Int, Int, Maybe Int)) -> Q Exp
genTrie 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 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 <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (\Exp
e -> (,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (forall (m :: * -> *). Quote m => String -> m Name
newName String
"rule") forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> 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 forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup a
k Map a a
m of
        Maybe a
Nothing -> forall a. HasCallStack => String -> a
error (String
"key not in map: " forall a. [a] -> [a] -> [a]
++ 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' = 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 = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Name -> Exp
VarE forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ 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 >> $(pure $ VarE $ fst $ ix branches 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
          | forall k a. Map k a -> Bool
M.null Map Word (Trie' (Maybe Int, Int, Maybe Int))
ts -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Name -> Exp
VarE forall a b. (a -> b) -> a -> b
$ forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ Map (Maybe Int) (Name, Exp)
branches forall k a. Ord k => Map k a -> k -> a
M.! Maybe Int
r
          | Bool
otherwise -> do
              ![(Word, Exp)]
next         <- (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse) Trie' (Maybe Int, Int, Maybe Int) -> Q Exp
go (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 forall a. Num a => a -> a -> a
+ Int
1)

              let cases :: Exp
cases = [Stmt] -> Exp
mkDoE 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"))
                         (forall a b. (a -> b) -> [a] -> [b]
map (\(Word
w, Exp
t) ->
                                 Pat -> Body -> [Dec] -> Match
Match (Lit -> Pat
LitP (Integer -> Lit
IntegerL (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
w)))
                                       (Exp -> Body
NormalB Exp
t)
                                       [])
                              [(Word, Exp)]
next
                          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    -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
cases
                Just Q Exp
alloc -> [| branch $alloc $(pure cases) $(fallback r 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 $(bytes ws) $(go t) $(fallback r n)|]
            Just Q Exp
alloc -> [| branch ($alloc >> $(bytesUnsafe ws)) $(go t) $(fallback r n) |]

  forall (m :: * -> *). Quote m => [m Dec] -> m Exp -> m Exp
letE
    (forall a b. (a -> b) -> [a] -> [b]
map (\(Name
x, Exp
rhs) -> forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
x) (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB (forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
rhs)) []) (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)

parseSwitch :: Q Exp -> Q ([(String, Exp)], Maybe Exp)
parseSwitch :: Q Exp -> Q ([(String, Exp)], Maybe Exp)
parseSwitch Q Exp
exp = Q Exp
exp forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  CaseE (UnboundVarE Name
_) []    -> forall a. HasCallStack => String -> a
error String
"switch: empty clause list"
  CaseE (UnboundVarE Name
_) [Match]
cases -> do
    (![Match]
cases, !Match
last) <- forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. [a] -> [a]
init [Match]
cases, forall a. [a] -> a
last [Match]
cases)
    ![(String, Exp)]
cases <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Match]
cases \case
      Match (LitP (StringL String
str)) (NormalB Exp
rhs) [] -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (String
str, Exp
rhs)
      Match
_ -> forall a. HasCallStack => String -> a
error String
"switch: expected a match clause on a string literal"
    (![(String, Exp)]
cases, !Maybe Exp
last) <- case Match
last of
      Match (LitP (StringL String
str)) (NormalB Exp
rhs) [] -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ([(String, Exp)]
cases forall a. [a] -> [a] -> [a]
++ [(String
str, Exp
rhs)], forall a. Maybe a
Nothing)
      Match Pat
WildP                (NormalB Exp
rhs) [] -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ([(String, Exp)]
cases, forall a. a -> Maybe a
Just Exp
rhs)
      Match
_ -> forall a. HasCallStack => String -> a
error String
"switch: expected a match clause on a string literal or a wildcard"
    forall (f :: * -> *) a. Applicative f => a -> f a
pure ([(String, Exp)]
cases, Maybe Exp
last)
  Exp
_ -> forall a. HasCallStack => String -> a
error String
"switch: expected a \"case _ of\" expression"

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) = forall a b. [(a, b)] -> ([a], [b])
unzip do
        (!Int
i, (!String
str, !Exp
rhs)) <- forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] [(String, Exp)]
cases
        case Maybe Exp
postAction of
          Maybe Exp
Nothing    -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ((forall a. a -> Maybe a
Just Int
i, Exp
rhs), (Int
i, String
str))
          Just !Exp
post -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ((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    = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ((forall a. Maybe a
Nothing, forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Name -> Exp
VarE 'failed) forall a. a -> a
id Maybe Exp
fallback) 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)