{-# OPTIONS_GHC -Wwarn #-}
{-# LANGUAGE NamedFieldPuns    #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Strict            #-}
module Language.Happy.Arbitrary where

import           Control.Applicative       ((<|>))
import           Control.Monad.Extra       (concatMapM)
import           Control.Monad.State.Lazy  (State)
import qualified Control.Monad.State.Lazy  as State
import           Data.Fix                  (foldFix)
import           Data.Map                  (Map)
import qualified Data.Map                  as Map
import           Data.Maybe                (fromJust)
import           Data.Text                 (Text)
import qualified Data.Text                 as Text
import           Debug.Trace               (traceM)
import           Language.Happy.Ast        (Node, NodeF (..))
import           Language.Happy.Lexer      (Lexeme, lexemeText)
import           Test.QuickCheck.Arbitrary (arbitrary)
import qualified Test.QuickCheck.Gen       as Gen
import           Test.QuickCheck.Gen       (Gen)

newtype Config token = Config
    { Config token -> Text -> token
parseToken :: Text -> token
    }

defConfig :: (Text -> token) -> Config token
defConfig :: (Text -> token) -> Config token
defConfig Text -> token
parseToken = Config :: forall token. (Text -> token) -> Config token
Config{Text -> token
parseToken :: Text -> token
parseToken :: Text -> token
parseToken}

genTokens :: Show token => Config token -> Text -> Node (Lexeme Text) -> Gen [token]
genTokens :: Config token -> Text -> Node (Lexeme Text) -> Gen [token]
genTokens Config token
cfg Text
start Node (Lexeme Text)
g = do
    [Int]
rng <- (Int -> Int) -> Gen [Int] -> Gen [Int]
forall a. (Int -> Int) -> Gen a -> Gen a
Gen.scale (Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
2) Gen [Int]
forall a. Arbitrary a => Gen a
arbitrary
    -- _ <- traceM $ "============================\nrng = " <> show rng
    [token] -> Gen [token]
forall (m :: * -> *) a. Monad m => a -> m a
return ([token] -> Gen [token]) -> [token] -> Gen [token]
forall a b. (a -> b) -> a -> b
$ case Text -> Map Text [[Text]] -> Maybe [[Text]]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
start Map Text [[Text]]
rules of
        Maybe [[Text]]
Nothing -> [Char] -> [token]
forall a. HasCallStack => [Char] -> a
error ([Char] -> [token]) -> [Char] -> [token]
forall a b. (a -> b) -> a -> b
$ [Char]
"no such rule: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Text -> [Char]
Text.unpack Text
start
        Just [[Text]]
r  -> State [Int] [token] -> [Int] -> [token]
forall s a. State s a -> s -> a
State.evalState (Config token
-> Map Text token
-> Map Text [[Text]]
-> [[Text]]
-> State [Int] [token]
forall token.
Show token =>
Config token
-> Map Text token
-> Map Text [[Text]]
-> [[Text]]
-> State [Int] [token]
expand Config token
cfg Map Text token
tokens Map Text [[Text]]
rules [[Text]]
r) [Int]
rng
  where
    tokens :: Map Text token
tokens = (NodeF (Lexeme Text) (Map Text token) -> Map Text token)
-> Node (Lexeme Text) -> Map Text token
forall (f :: * -> *) a. Functor f => (f a -> a) -> Fix f -> a
foldFix (Config token
-> NodeF (Lexeme Text) (Map Text token) -> Map Text token
forall token.
Config token
-> NodeF (Lexeme Text) (Map Text token) -> Map Text token
terminals Config token
cfg) Node (Lexeme Text)
g
    rules :: Map Text [[Text]]
rules = (NodeF (Lexeme Text) (Map Text [[Text]]) -> Map Text [[Text]])
-> Node (Lexeme Text) -> Map Text [[Text]]
forall (f :: * -> *) a. Functor f => (f a -> a) -> Fix f -> a
foldFix NodeF (Lexeme Text) (Map Text [[Text]]) -> Map Text [[Text]]
nonterminals Node (Lexeme Text)
g

expand :: Show token => Config token -> Map Text token -> Map Text [[Text]] -> [[Text]] -> State [Int] [token]
expand :: Config token
-> Map Text token
-> Map Text [[Text]]
-> [[Text]]
-> State [Int] [token]
expand Config token
cfg Map Text token
tokens Map Text [[Text]]
rules [[Text]]
nonterm = do
    -- _ <- traceM $ show rng <> ": selecting from " <> show nonterm
    [Text]
rule <- [[Text]] -> State [Int] [Text]
forall a. [a] -> State [Int] a
select [[Text]]
nonterm
    [token]
res <- (Text -> State [Int] [token]) -> [Text] -> State [Int] [token]
forall (m :: * -> *) a b. Monad m => (a -> m [b]) -> [a] -> m [b]
concatMapM (Config token
-> Map Text token
-> Map Text [[Text]]
-> Either token [[Text]]
-> State [Int] [token]
forall token.
Show token =>
Config token
-> Map Text token
-> Map Text [[Text]]
-> Either token [[Text]]
-> State [Int] [token]
continue Config token
cfg Map Text token
tokens Map Text [[Text]]
rules (Either token [[Text]] -> State [Int] [token])
-> (Text -> Either token [[Text]]) -> Text -> State [Int] [token]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Text token
-> Map Text [[Text]] -> Text -> Either token [[Text]]
forall token.
Show token =>
Map Text token
-> Map Text [[Text]] -> Text -> Either token [[Text]]
resolve Map Text token
tokens Map Text [[Text]]
rules) [Text]
rule
    -- _ <- traceM $ show rng <> ": result: " <> show res
    [token] -> State [Int] [token]
forall (m :: * -> *) a. Monad m => a -> m a
return [token]
res

select :: [a] -> State [Int] a
select :: [a] -> State [Int] a
select [] = [Char] -> State [Int] a
forall a. HasCallStack => [Char] -> a
error [Char]
"nope"
select nonterm :: [a]
nonterm@(a
rule:[a]
_) = do
    [Int]
rng <- StateT [Int] Identity [Int]
forall s (m :: * -> *). MonadState s m => m s
State.get
    case [Int]
rng of
        [] -> a -> State [Int] a
forall (m :: * -> *) a. Monad m => a -> m a
return a
rule
        (Int
i:[Int]
is) -> do
            [Int] -> StateT [Int] Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
State.put [Int]
is
            a -> State [Int] a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> State [Int] a) -> a -> State [Int] a
forall a b. (a -> b) -> a -> b
$ [a]
nonterm [a] -> Int -> a
forall a. [a] -> Int -> a
!! (Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
nonterm)

continue :: Show token => Config token -> Map Text token -> Map Text [[Text]] -> Either token [[Text]] -> State [Int] [token]
continue :: Config token
-> Map Text token
-> Map Text [[Text]]
-> Either token [[Text]]
-> State [Int] [token]
continue Config token
_ Map Text token
_ Map Text [[Text]]
_ (Left token
token)            = [token] -> State [Int] [token]
forall (m :: * -> *) a. Monad m => a -> m a
return [token
token]
continue Config token
cfg Map Text token
tokens Map Text [[Text]]
rules (Right [[Text]]
rule) = Config token
-> Map Text token
-> Map Text [[Text]]
-> [[Text]]
-> State [Int] [token]
forall token.
Show token =>
Config token
-> Map Text token
-> Map Text [[Text]]
-> [[Text]]
-> State [Int] [token]
expand Config token
cfg Map Text token
tokens Map Text [[Text]]
rules [[Text]]
rule

resolve :: Show token => Map Text token -> Map Text [[Text]] -> Text -> Either token [[Text]]
resolve :: Map Text token
-> Map Text [[Text]] -> Text -> Either token [[Text]]
resolve Map Text token
tokens Map Text [[Text]]
rules Text
sym =
    Maybe (Either token [[Text]]) -> Either token [[Text]]
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe (Either token [[Text]]) -> Either token [[Text]])
-> Maybe (Either token [[Text]]) -> Either token [[Text]]
forall a b. (a -> b) -> a -> b
$ (token -> Either token [[Text]]
forall a b. a -> Either a b
Left (token -> Either token [[Text]])
-> Maybe token -> Maybe (Either token [[Text]])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> Map Text token -> Maybe token
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
sym Map Text token
tokens) Maybe (Either token [[Text]])
-> Maybe (Either token [[Text]]) -> Maybe (Either token [[Text]])
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ([[Text]] -> Either token [[Text]]
forall a b. b -> Either a b
Right ([[Text]] -> Either token [[Text]])
-> Maybe [[Text]] -> Maybe (Either token [[Text]])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> Map Text [[Text]] -> Maybe [[Text]]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
sym Map Text [[Text]]
rules)

terminals :: Config token -> NodeF (Lexeme Text) (Map Text token) -> Map Text token
terminals :: Config token
-> NodeF (Lexeme Text) (Map Text token) -> Map Text token
terminals Config{Text -> token
parseToken :: Text -> token
parseToken :: forall token. Config token -> Text -> token
parseToken} NodeF (Lexeme Text) (Map Text token)
node = case NodeF (Lexeme Text) (Map Text token)
node of
    Token Lexeme Text
k Lexeme Text
v -> Text -> token -> Map Text token
forall k a. k -> a -> Map k a
Map.singleton (Lexeme Text -> Text
forall text. Lexeme text -> text
lexemeText Lexeme Text
k) (Text -> token
parseToken (Text -> token) -> Text -> token
forall a b. (a -> b) -> a -> b
$ Lexeme Text -> Text
forall text. Lexeme text -> text
lexemeText Lexeme Text
v)
    NodeF (Lexeme Text) (Map Text token)
n         -> NodeF (Lexeme Text) (Map Text token) -> Map Text token
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
Map.unions NodeF (Lexeme Text) (Map Text token)
n

nonterminals :: NodeF (Lexeme Text) (Map Text [[Text]]) -> Map Text [[Text]]
nonterminals :: NodeF (Lexeme Text) (Map Text [[Text]]) -> Map Text [[Text]]
nonterminals NodeF (Lexeme Text) (Map Text [[Text]])
node = case NodeF (Lexeme Text) (Map Text [[Text]])
node of
    RuleLine [Lexeme Text]
syms Lexeme Text
_     -> Text -> [[Text]] -> Map Text [[Text]]
forall k a. k -> a -> Map k a
Map.singleton Text
"" [(Lexeme Text -> Text) -> [Lexeme Text] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Lexeme Text -> Text
forall text. Lexeme text -> text
lexemeText [Lexeme Text]
syms]
    RuleDefn Lexeme Text
name [Map Text [[Text]]]
rules -> Text -> [[Text]] -> Map Text [[Text]]
forall k a. k -> a -> Map k a
Map.singleton (Lexeme Text -> Text
forall text. Lexeme text -> text
lexemeText Lexeme Text
name) ([Map Text [[Text]]] -> [[Text]]
forall k a. [Map k [a]] -> [a]
merge [Map Text [[Text]]]
rules)
    NodeF (Lexeme Text) (Map Text [[Text]])
n                   -> ([[Text]] -> [[Text]] -> [[Text]])
-> NodeF (Lexeme Text) (Map Text [[Text]]) -> Map Text [[Text]]
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
Map.unionsWith [[Text]] -> [[Text]] -> [[Text]]
forall a. [a] -> [a] -> [a]
(++) NodeF (Lexeme Text) (Map Text [[Text]])
n
  where
    merge :: [Map k [a]] -> [a]
merge = [[a]] -> [a]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[a]] -> [a]) -> ([Map k [a]] -> [[a]]) -> [Map k [a]] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map k [a] -> [[a]]) -> [Map k [a]] -> [[a]]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Map k [a] -> [[a]]
forall k a. Map k a -> [a]
Map.elems