{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE QualifiedDo #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -Wno-orphans #-}

module Language.Haskell.TH.Natural.Syntax.Expr.Simple.Untyped (
    -- * Builder
    newExpr,
    SimpleExprBuilder,

    -- * State
    module Language.Haskell.TH.Natural.Syntax.Expr.Simple.State,

    -- * Operations
    arg,

    -- * Reexport
    module Language.Haskell.TH.Natural.Syntax.Expr.Untyped,
    module Language.Haskell.TH.Natural.Syntax.Builder.Monad,
) where

import Control.Lens (views, (?=), (^.), (|>=))
import Control.Monad
import Data.Bifunctor
import Data.Functor ((<&>))
import Data.List (partition)
import Data.Maybe
import qualified Language.Haskell.TH as TH
import Language.Haskell.TH.Gen
import Language.Haskell.TH.Natural.Syntax.Builder hiding (fail, (>>=))
import Language.Haskell.TH.Natural.Syntax.Builder.Monad
import Language.Haskell.TH.Natural.Syntax.Expr.Internal
import Language.Haskell.TH.Natural.Syntax.Expr.Simple.State
import Language.Haskell.TH.Natural.Syntax.Expr.Untyped
import Language.Haskell.TH.Syntax.ExtractedCons (LamE (..))
import Prelude hiding ((>>=))

type SimpleExprBuilder = Builder SimpleExprBuilderState

-- | Builds a lambda/let expression
newExpr :: SimpleExprBuilder step Ready () -> TH.Q LamE
newExpr :: forall (step :: BuilderStep).
SimpleExprBuilder step 'Ready () -> Q LamE
newExpr = BaseBuilder Q SimpleExprBuilderState step 'Ready () -> Q LamE
BaseBuilder Q SimpleExprBuilderState step 'Ready ()
-> Definition SimpleExprBuilderState
forall st (step :: BuilderStep).
IsExprBuilder st =>
Builder st step 'Ready () -> Definition st
forall (step :: BuilderStep).
Builder SimpleExprBuilderState step 'Ready ()
-> Definition SimpleExprBuilderState
runExprBuilder

-- | Add an argument to the lambda
arg :: SimpleExprBuilder curr curr TH.Exp
arg :: forall {k} (curr :: k). SimpleExprBuilder curr curr Exp
arg = do
    nextArgName <- Q Name -> BaseBuilder Q SimpleExprBuilderState curr curr Name
forall {k} (m :: * -> *) a s (step :: k).
Monad m =>
m a -> BaseBuilder m s step step a
liftB (Q Name -> BaseBuilder Q SimpleExprBuilderState curr curr Name)
-> Q Name -> BaseBuilder Q SimpleExprBuilderState curr curr Name
forall a b. (a -> b) -> a -> b
$ String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
TH.newName String
"arg"
    argNames |>= nextArgName
    return $ TH.VarE nextArgName

instance IsExprBuilder SimpleExprBuilderState where
    type Definition SimpleExprBuilderState = TH.Q LamE
    returns :: forall b (step :: BuilderStep).
GenExpr b =>
b -> Builder SimpleExprBuilderState step 'Ready ()
returns b
q = BaseBuilder Q SimpleExprBuilderState (ZonkAny 1) (ZonkAny 1) ()
-> BaseBuilder Q SimpleExprBuilderState step 'Ready ()
forall {k1} {k2} (prev' :: k1) (curr' :: k1) (prev :: k2)
       (curr :: k2) (m :: * -> *) s a.
BaseBuilder m s prev curr a -> BaseBuilder m s prev' curr' a
unsafeCastStep (BaseBuilder Q SimpleExprBuilderState (ZonkAny 1) (ZonkAny 1) ()
 -> BaseBuilder Q SimpleExprBuilderState step 'Ready ())
-> BaseBuilder Q SimpleExprBuilderState (ZonkAny 1) (ZonkAny 1) ()
-> BaseBuilder Q SimpleExprBuilderState step 'Ready ()
forall a b. (a -> b) -> a -> b
$ do
        expr <- Q Exp
-> BaseBuilder Q SimpleExprBuilderState (ZonkAny 1) (ZonkAny 1) Exp
forall {k} (m :: * -> *) a s (step :: k).
Monad m =>
m a -> BaseBuilder m s step step a
liftB (Q Exp
 -> BaseBuilder
      Q SimpleExprBuilderState (ZonkAny 1) (ZonkAny 1) Exp)
-> Q Exp
-> BaseBuilder Q SimpleExprBuilderState (ZonkAny 1) (ZonkAny 1) Exp
forall a b. (a -> b) -> a -> b
$ b -> Q Exp
forall a. GenExpr a => a -> Q Exp
genExpr b
q
        returnedExp ?= expr

    addDeconstruct :: forall (step :: BuilderStep).
Deconstruct -> Builder SimpleExprBuilderState step 'Empty ()
addDeconstruct Deconstruct
d = BaseBuilder Q SimpleExprBuilderState step step ()
-> BaseBuilder Q SimpleExprBuilderState step 'Empty ()
forall {k} (m :: * -> *) s (step :: k) a (next :: k).
BaseBuilder m s step step a -> BaseBuilder m s step next a
impure (BaseBuilder Q SimpleExprBuilderState step step ()
 -> BaseBuilder Q SimpleExprBuilderState step 'Empty ())
-> BaseBuilder Q SimpleExprBuilderState step step ()
-> BaseBuilder Q SimpleExprBuilderState step 'Empty ()
forall a b. (a -> b) -> a -> b
$ ([Deconstruct] -> Identity [Deconstruct])
-> SimpleExprBuilderState -> Identity SimpleExprBuilderState
Lens' SimpleExprBuilderState [Deconstruct]
deconstructs (([Deconstruct] -> Identity [Deconstruct])
 -> SimpleExprBuilderState -> Identity SimpleExprBuilderState)
-> Deconstruct -> BaseBuilder Q SimpleExprBuilderState step step ()
forall s (m :: * -> *) b a.
(MonadState s m, Snoc b b a a) =>
ASetter s s b b -> a -> m ()
|>= Deconstruct
d

    addLet :: forall (step :: BuilderStep).
Binding -> Builder SimpleExprBuilderState step 'Empty ()
addLet Binding
l = BaseBuilder Q SimpleExprBuilderState step step ()
-> BaseBuilder Q SimpleExprBuilderState step 'Empty ()
forall {k} (m :: * -> *) s (step :: k) a (next :: k).
BaseBuilder m s step step a -> BaseBuilder m s step next a
impure (BaseBuilder Q SimpleExprBuilderState step step ()
 -> BaseBuilder Q SimpleExprBuilderState step 'Empty ())
-> BaseBuilder Q SimpleExprBuilderState step step ()
-> BaseBuilder Q SimpleExprBuilderState step 'Empty ()
forall a b. (a -> b) -> a -> b
$ ([Binding] -> Identity [Binding])
-> SimpleExprBuilderState -> Identity SimpleExprBuilderState
Lens' SimpleExprBuilderState [Binding]
lets (([Binding] -> Identity [Binding])
 -> SimpleExprBuilderState -> Identity SimpleExprBuilderState)
-> Binding -> BaseBuilder Q SimpleExprBuilderState step step ()
forall s (m :: * -> *) b a.
(MonadState s m, Snoc b b a a) =>
ASetter s s b b -> a -> m ()
|>= Binding
l

    letCount :: forall {k} (step :: k).
Builder SimpleExprBuilderState step step Int
letCount = LensLike' (Const Int) SimpleExprBuilderState [Binding]
-> ([Binding] -> Int)
-> BaseBuilder Q SimpleExprBuilderState step step Int
forall s (m :: * -> *) r a.
MonadReader s m =>
LensLike' (Const r) s a -> (a -> r) -> m r
views LensLike' (Const Int) SimpleExprBuilderState [Binding]
Lens' SimpleExprBuilderState [Binding]
lets [Binding] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length

    runExprBuilder :: forall (step :: BuilderStep).
Builder SimpleExprBuilderState step 'Ready ()
-> Definition SimpleExprBuilderState
runExprBuilder Builder SimpleExprBuilderState step 'Ready ()
b = do
        st <- Builder SimpleExprBuilderState step 'Ready ()
-> SimpleExprBuilderState -> Q SimpleExprBuilderState
forall {k} (m :: * -> *) s (step :: k) (end :: k).
Monad m =>
BaseBuilder m s step end () -> s -> m s
runBaseBuilder Builder SimpleExprBuilderState step 'Ready ()
b ([Name]
-> [Binding]
-> [Deconstruct]
-> Maybe Exp
-> SimpleExprBuilderState
MkEBS [] [] [] Maybe Exp
forall a. Maybe a
Nothing)
        (argsPat, decs) <- _compileSimpleExpr st
        let letOrId = case [Dec]
decs of
                [] -> Exp -> Exp
forall a. a -> a
id
                [Dec]
_ -> [Dec] -> Exp -> Exp
TH.LetE [Dec]
decs
        resExp <- case st ^. returnedExp of
            Maybe Exp
Nothing -> String -> Q Exp
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
Prelude.fail String
"Missing returned expression"
            Just Exp
e -> Exp -> Q Exp
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp
e
        return $ MkLamE argsPat $ letOrId resExp

-- | Merges deconstructions of common expressions together, and returns a pair where:
--
-- - The first element is the patterns for the args (if any), deconstructed or not according to the decons
-- - The list of declarations (i.e. lets and deconstructs) to pass to 'LetE'
_compileSimpleExpr :: (MonadFail m) => SimpleExprBuilderState -> m ([TH.Pat], [TH.Dec])
_compileSimpleExpr :: forall (m :: * -> *).
MonadFail m =>
SimpleExprBuilderState -> m ([Pat], [Dec])
_compileSimpleExpr SimpleExprBuilderState
st = do
    mergedDecons <- [Deconstruct] -> m [Deconstruct]
forall (m :: * -> *).
MonadFail m =>
[Deconstruct] -> m [Deconstruct]
mergeDeconstructs ([Deconstruct] -> m [Deconstruct])
-> [Deconstruct] -> m [Deconstruct]
forall a b. (a -> b) -> a -> b
$ SimpleExprBuilderState
st SimpleExprBuilderState
-> Getting [Deconstruct] SimpleExprBuilderState [Deconstruct]
-> [Deconstruct]
forall s a. s -> Getting a s a -> a
^. Getting [Deconstruct] SimpleExprBuilderState [Deconstruct]
Lens' SimpleExprBuilderState [Deconstruct]
deconstructs
    let (argPats, decons') =
            foldr
                (\Name
argName ([Pat]
argsPats, [Deconstruct]
decons) -> (Pat -> [Pat]) -> (Pat, [Deconstruct]) -> ([Pat], [Deconstruct])
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (Pat -> [Pat] -> [Pat]
forall a. a -> [a] -> [a]
: [Pat]
argsPats) ((Pat, [Deconstruct]) -> ([Pat], [Deconstruct]))
-> (Pat, [Deconstruct]) -> ([Pat], [Deconstruct])
forall a b. (a -> b) -> a -> b
$ Name -> [Deconstruct] -> (Pat, [Deconstruct])
argToPat Name
argName [Deconstruct]
decons)
                ([], mergedDecons)
                (st ^. argNames)
        bindDecs = SimpleExprBuilderState
st SimpleExprBuilderState
-> Getting [Binding] SimpleExprBuilderState [Binding] -> [Binding]
forall s a. s -> Getting a s a -> a
^. Getting [Binding] SimpleExprBuilderState [Binding]
Lens' SimpleExprBuilderState [Binding]
lets [Binding] -> (Binding -> Dec) -> [Dec]
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> Binding -> Dec
bindingToDec
        deconDecs = [Deconstruct]
decons' [Deconstruct] -> (Deconstruct -> Dec) -> [Dec]
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> Deconstruct -> Dec
deconstructToDec
    return (argPats, bindDecs ++ deconDecs)
  where
    argToPat :: TH.Name -> [Deconstruct] -> (TH.Pat, [Deconstruct])
    argToPat :: Name -> [Deconstruct] -> (Pat, [Deconstruct])
argToPat Name
n [Deconstruct]
decons = case (Deconstruct -> Bool)
-> [Deconstruct] -> ([Deconstruct], [Deconstruct])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition (\Deconstruct
d -> Deconstruct -> Exp
_src Deconstruct
d Exp -> Exp -> Bool
forall a. Eq a => a -> a -> Bool
== Exp
nExp) [Deconstruct]
decons of
        ([], [Deconstruct]
_) -> (Name -> Pat
TH.VarP Name
n, [Deconstruct]
decons)
        -- Only considering one matching decons, this function should be called using merged decons.
        (Deconstruct
match : [Deconstruct]
decons', [Deconstruct]
decons'') -> (Deconstruct -> Pat
deconstructToPat Deconstruct
match, [Deconstruct]
decons' [Deconstruct] -> [Deconstruct] -> [Deconstruct]
forall a. [a] -> [a] -> [a]
++ [Deconstruct]
decons'')
      where
        nExp :: Exp
nExp = Name -> Exp
TH.VarE Name
n