{-# OPTIONS_GHC -Wno-orphans #-}

module Language.Haskell.TH.Natural.Syntax.Expr.Do.Untyped (
    -- * Types
    newDo,
    newQualifiedDo,
    DoExprDefinition,
    DoExprBuilder,

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

    -- * Functions
    stmt,
    bind,
    strictBind,
    bind_,

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

import Control.Lens hiding (Empty)
import Control.Monad (foldM)
import Data.Bifunctor
import qualified Language.Haskell.TH as TH
import Language.Haskell.TH.Gen (GenExpr (genExpr))
import Language.Haskell.TH.Natural.Syntax.Builder
import Language.Haskell.TH.Natural.Syntax.Builder.Monad
import Language.Haskell.TH.Natural.Syntax.Expr.Do.State
import Language.Haskell.TH.Natural.Syntax.Expr.Internal
import Language.Haskell.TH.Natural.Syntax.Expr.Untyped
import Language.Haskell.TH.Syntax (ModName (..), nameBase)
import Language.Haskell.TH.Syntax.ExtractedCons hiding (expr)

type DoExprDefinition = TH.Q DoE
type DoExprBuilder = Builder DoExprBuilderState

-- | Builds a do-expression
newDo :: DoExprBuilder step Ready () -> DoExprDefinition
newDo :: forall (step :: BuilderStep).
DoExprBuilder step 'Ready () -> DoExprDefinition
newDo = BaseBuilder Q DoExprBuilderState step 'Ready () -> DoExprDefinition
BaseBuilder Q DoExprBuilderState step 'Ready ()
-> Definition DoExprBuilderState
forall st (step :: BuilderStep).
IsExprBuilder st =>
Builder st step 'Ready () -> Definition st
forall (step :: BuilderStep).
Builder DoExprBuilderState step 'Ready ()
-> Definition DoExprBuilderState
runExprBuilder

-- | Same as 'newDo', but allow qualifying the do-expression
newQualifiedDo :: TH.Name -> DoExprBuilder step Ready () -> DoExprDefinition
newQualifiedDo :: forall (step :: BuilderStep).
Name -> DoExprBuilder step 'Ready () -> DoExprDefinition
newQualifiedDo Name
modN DoExprBuilder step 'Ready ()
builder = do
    doE <- DoExprBuilder step 'Ready () -> Definition DoExprBuilderState
forall st (step :: BuilderStep).
IsExprBuilder st =>
Builder st step 'Ready () -> Definition st
forall (step :: BuilderStep).
Builder DoExprBuilderState step 'Ready ()
-> Definition DoExprBuilderState
runExprBuilder DoExprBuilder step 'Ready ()
builder
    return $ over modName (const $ Just $ ModName $ nameBase modN) doE

-- | Add a statement to the do-expression
stmt :: (GenExpr b) => b -> DoExprBuilder step Ready ()
stmt :: forall b (step :: BuilderStep).
GenExpr b =>
b -> DoExprBuilder step 'Ready ()
stmt = b -> Builder DoExprBuilderState step 'Ready ()
forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
b -> Builder st step 'Ready ()
forall b (step :: BuilderStep).
GenExpr b =>
b -> DoExprBuilder step 'Ready ()
returns

-- | Bind a value using the backward arrow
bind :: (GenExpr b) => b -> DoExprBuilder step Empty TH.Exp
bind :: forall b (step :: BuilderStep).
GenExpr b =>
b -> DoExprBuilder step 'Empty Exp
bind = Bool -> b -> DoExprBuilder step 'Empty Exp
forall b (step :: BuilderStep).
GenExpr b =>
Bool -> b -> DoExprBuilder step 'Empty Exp
bind_ Bool
False

-- | Same as 'bind', but uses a bang pattern
strictBind :: (GenExpr b) => b -> DoExprBuilder step Empty TH.Exp
strictBind :: forall b (step :: BuilderStep).
GenExpr b =>
b -> DoExprBuilder step 'Empty Exp
strictBind = Bool -> b -> DoExprBuilder step 'Empty Exp
forall b (step :: BuilderStep).
GenExpr b =>
Bool -> b -> DoExprBuilder step 'Empty Exp
bind_ Bool
True

bind_ :: (GenExpr b) => Bool -> b -> DoExprBuilder step Empty TH.Exp
bind_ :: forall b (step :: BuilderStep).
GenExpr b =>
Bool -> b -> DoExprBuilder step 'Empty Exp
bind_ Bool
s b
q = BaseBuilder Q DoExprBuilderState step step Exp
-> BaseBuilder Q DoExprBuilderState step 'Empty Exp
forall {k} (m :: * -> *) s (step :: k) a (next :: k).
BaseBuilder m s step step a -> BaseBuilder m s step next a
impure (BaseBuilder Q DoExprBuilderState step step Exp
 -> BaseBuilder Q DoExprBuilderState step 'Empty Exp)
-> BaseBuilder Q DoExprBuilderState step step Exp
-> BaseBuilder Q DoExprBuilderState step 'Empty Exp
forall a b. (a -> b) -> a -> b
$ do
    stepCount <-
        LensLike' (Const Int) DoExprBuilderState [DoExprStep]
-> ([DoExprStep] -> Int)
-> BaseBuilder Q DoExprBuilderState step step Int
forall s (m :: * -> *) r a.
MonadReader s m =>
LensLike' (Const r) s a -> (a -> r) -> m r
views LensLike' (Const Int) DoExprBuilderState [DoExprStep]
Iso' DoExprBuilderState [DoExprStep]
steps (([DoExprStep] -> Int)
 -> BaseBuilder Q DoExprBuilderState step step Int)
-> ([DoExprStep] -> Int)
-> BaseBuilder Q DoExprBuilderState step step Int
forall a b. (a -> b) -> a -> b
$
            [DoExprStep] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length
                ([DoExprStep] -> Int)
-> ([DoExprStep] -> [DoExprStep]) -> [DoExprStep] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (DoExprStep -> Bool) -> [DoExprStep] -> [DoExprStep]
forall a. (a -> Bool) -> [a] -> [a]
filter
                    ( \case
                        Bind Binding
_ -> Bool
True
                        DoExprStep
_ -> Bool
False
                    )
    varName <- liftB $ TH.newName ("var" ++ show stepCount)
    e <- liftB $ genExpr q
    steps <|= Bind (MkBind varName e s)
    return $ TH.VarE varName

instance IsExprBuilder DoExprBuilderState where
    type Definition DoExprBuilderState = DoExprDefinition
    returns :: forall b (step :: BuilderStep).
GenExpr b =>
b -> DoExprBuilder step 'Ready ()
returns b
q = BaseBuilder Q DoExprBuilderState (ZonkAny 1) (ZonkAny 1) ()
-> BaseBuilder Q DoExprBuilderState 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 DoExprBuilderState (ZonkAny 1) (ZonkAny 1) ()
 -> BaseBuilder Q DoExprBuilderState step 'Ready ())
-> BaseBuilder Q DoExprBuilderState (ZonkAny 1) (ZonkAny 1) ()
-> BaseBuilder Q DoExprBuilderState step 'Ready ()
forall a b. (a -> b) -> a -> b
$ do
        e <- Q Exp
-> BaseBuilder Q DoExprBuilderState (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 DoExprBuilderState (ZonkAny 1) (ZonkAny 1) Exp)
-> Q Exp
-> BaseBuilder Q DoExprBuilderState (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
        steps <|= Stmt e

    letCount :: forall {k} (step :: k). Builder DoExprBuilderState step step Int
letCount =
        LensLike' (Const Int) DoExprBuilderState [DoExprStep]
-> ([DoExprStep] -> Int)
-> BaseBuilder Q DoExprBuilderState step step Int
forall s (m :: * -> *) r a.
MonadReader s m =>
LensLike' (Const r) s a -> (a -> r) -> m r
views LensLike' (Const Int) DoExprBuilderState [DoExprStep]
Iso' DoExprBuilderState [DoExprStep]
steps (([DoExprStep] -> Int)
 -> BaseBuilder Q DoExprBuilderState step step Int)
-> ([DoExprStep] -> Int)
-> BaseBuilder Q DoExprBuilderState step step Int
forall a b. (a -> b) -> a -> b
$
            [DoExprStep] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length
                ([DoExprStep] -> Int)
-> ([DoExprStep] -> [DoExprStep]) -> [DoExprStep] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (DoExprStep -> Bool) -> [DoExprStep] -> [DoExprStep]
forall a. (a -> Bool) -> [a] -> [a]
filter
                    ( \case
                        Let Binding
_ -> Bool
True
                        DoExprStep
_ -> Bool
False
                    )
    addDeconstruct :: forall (step :: BuilderStep).
Deconstruct -> Builder DoExprBuilderState step 'Empty ()
addDeconstruct Deconstruct
d = BaseBuilder Q DoExprBuilderState step step ()
-> BaseBuilder Q DoExprBuilderState 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 DoExprBuilderState step step ()
 -> BaseBuilder Q DoExprBuilderState step 'Empty ())
-> BaseBuilder Q DoExprBuilderState step step ()
-> BaseBuilder Q DoExprBuilderState step 'Empty ()
forall a b. (a -> b) -> a -> b
$ ([DoExprStep] -> Identity [DoExprStep])
-> DoExprBuilderState -> Identity DoExprBuilderState
Iso' DoExprBuilderState [DoExprStep]
steps (([DoExprStep] -> Identity [DoExprStep])
 -> DoExprBuilderState -> Identity DoExprBuilderState)
-> DoExprStep -> BaseBuilder Q DoExprBuilderState step step ()
forall s (m :: * -> *) b a.
(MonadState s m, Cons b b a a) =>
ASetter s s b b -> a -> m ()
<|= Deconstruct -> DoExprStep
Decons Deconstruct
d

    addLet :: forall (step :: BuilderStep).
Binding -> Builder DoExprBuilderState step 'Empty ()
addLet Binding
b = BaseBuilder Q DoExprBuilderState step step ()
-> BaseBuilder Q DoExprBuilderState 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 DoExprBuilderState step step ()
 -> BaseBuilder Q DoExprBuilderState step 'Empty ())
-> BaseBuilder Q DoExprBuilderState step step ()
-> BaseBuilder Q DoExprBuilderState step 'Empty ()
forall a b. (a -> b) -> a -> b
$ ([DoExprStep] -> Identity [DoExprStep])
-> DoExprBuilderState -> Identity DoExprBuilderState
Iso' DoExprBuilderState [DoExprStep]
steps (([DoExprStep] -> Identity [DoExprStep])
 -> DoExprBuilderState -> Identity DoExprBuilderState)
-> DoExprStep -> BaseBuilder Q DoExprBuilderState step step ()
forall s (m :: * -> *) b a.
(MonadState s m, Cons b b a a) =>
ASetter s s b b -> a -> m ()
<|= Binding -> DoExprStep
Let Binding
b

    runExprBuilder :: forall (step :: BuilderStep).
Builder DoExprBuilderState step 'Ready ()
-> Definition DoExprBuilderState
runExprBuilder Builder DoExprBuilderState step 'Ready ()
b = do
        MkDoEBS doSteps <- Builder DoExprBuilderState step 'Ready ()
-> DoExprBuilderState -> Q DoExprBuilderState
forall {k} (m :: * -> *) s (step :: k) (end :: k).
Monad m =>
BaseBuilder m s step end () -> s -> m s
runBaseBuilder Builder DoExprBuilderState step 'Ready ()
b ([DoExprStep] -> DoExprBuilderState
MkDoEBS [])
        mergedSteps <- _compileDoExpr $ reverse doSteps
        return $ MkDoE Nothing $ fmap stepToStmt mergedSteps
      where
        stepToStmt :: DoExprStep -> Stmt
stepToStmt = \case
            Bind (MkBind Name
n Exp
e Bool
s) -> Pat -> Exp -> Stmt
TH.BindS ((if Bool
s then Pat -> Pat
TH.BangP else Pat -> Pat
forall a. a -> a
id) (Pat -> Pat) -> Pat -> Pat
forall a b. (a -> b) -> a -> b
$ Name -> Pat
TH.VarP Name
n) Exp
e
            Stmt Exp
e -> Exp -> Stmt
TH.NoBindS Exp
e
            Let Binding
let_ -> [Dec] -> Stmt
TH.LetS [Binding -> Dec
bindingToDec Binding
let_]
            Decons Deconstruct
decons -> [Dec] -> Stmt
TH.LetS [Deconstruct -> Dec
deconstructToDec Deconstruct
decons]

-- | Merges deconstructions of common expressions together
_compileDoExpr :: (MonadFail m) => [DoExprStep] -> m [DoExprStep]
_compileDoExpr :: forall (m :: * -> *). MonadFail m => [DoExprStep] -> m [DoExprStep]
_compileDoExpr = [DoExprStep] -> m [DoExprStep]
forall (m :: * -> *). MonadFail m => [DoExprStep] -> m [DoExprStep]
mergeDeconsSteps
  where
    -- TODO: Merge Bind and decons, but 'Binding' uses a name for the pattern
    mergeDeconsSteps :: [DoExprStep] -> f [DoExprStep]
mergeDeconsSteps [] = [DoExprStep] -> f [DoExprStep]
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
    mergeDeconsSteps (Decons Deconstruct
d : [DoExprStep]
rest) = case Exp -> [DoExprStep] -> ([Deconstruct], [DoExprStep])
partitionMatchingDecons (Deconstruct -> Exp
_src Deconstruct
d) [DoExprStep]
rest of
        ([], [DoExprStep]
_) -> (Deconstruct -> DoExprStep
Decons Deconstruct
d DoExprStep -> [DoExprStep] -> [DoExprStep]
forall a. a -> [a] -> [a]
:) ([DoExprStep] -> [DoExprStep]) -> f [DoExprStep] -> f [DoExprStep]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [DoExprStep] -> f [DoExprStep]
mergeDeconsSteps [DoExprStep]
rest
        ([Deconstruct]
ds, [DoExprStep]
rest') -> do
            d' <- Deconstruct -> DoExprStep
Decons (Deconstruct -> DoExprStep) -> f Deconstruct -> f DoExprStep
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Deconstruct -> Deconstruct -> f Deconstruct)
-> Deconstruct -> [Deconstruct] -> f Deconstruct
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM Deconstruct -> Deconstruct -> f Deconstruct
forall (m :: * -> *).
MonadFail m =>
Deconstruct -> Deconstruct -> m Deconstruct
mergeDeconstruct Deconstruct
d [Deconstruct]
ds
            (d' :) <$> mergeDeconsSteps rest'
    mergeDeconsSteps (DoExprStep
s : [DoExprStep]
rest) = (DoExprStep
s DoExprStep -> [DoExprStep] -> [DoExprStep]
forall a. a -> [a] -> [a]
:) ([DoExprStep] -> [DoExprStep]) -> f [DoExprStep] -> f [DoExprStep]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [DoExprStep] -> f [DoExprStep]
mergeDeconsSteps [DoExprStep]
rest
    partitionMatchingDecons :: TH.Exp -> [DoExprStep] -> ([Deconstruct], [DoExprStep])
    partitionMatchingDecons :: Exp -> [DoExprStep] -> ([Deconstruct], [DoExprStep])
partitionMatchingDecons Exp
_ [] = ([], [])
    partitionMatchingDecons Exp
e (Decons Deconstruct
d : [DoExprStep]
r) =
        if Deconstruct -> Exp
_src Deconstruct
d Exp -> Exp -> Bool
forall a. Eq a => a -> a -> Bool
== Exp
e
            then ([Deconstruct] -> [Deconstruct])
-> ([Deconstruct], [DoExprStep]) -> ([Deconstruct], [DoExprStep])
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 (Deconstruct
d Deconstruct -> [Deconstruct] -> [Deconstruct]
forall a. a -> [a] -> [a]
:) (([Deconstruct], [DoExprStep]) -> ([Deconstruct], [DoExprStep]))
-> ([Deconstruct], [DoExprStep]) -> ([Deconstruct], [DoExprStep])
forall a b. (a -> b) -> a -> b
$ Exp -> [DoExprStep] -> ([Deconstruct], [DoExprStep])
partitionMatchingDecons Exp
e [DoExprStep]
r
            else ([DoExprStep] -> [DoExprStep])
-> ([Deconstruct], [DoExprStep]) -> ([Deconstruct], [DoExprStep])
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (Deconstruct -> DoExprStep
Decons Deconstruct
d DoExprStep -> [DoExprStep] -> [DoExprStep]
forall a. a -> [a] -> [a]
:) (([Deconstruct], [DoExprStep]) -> ([Deconstruct], [DoExprStep]))
-> ([Deconstruct], [DoExprStep]) -> ([Deconstruct], [DoExprStep])
forall a b. (a -> b) -> a -> b
$ Exp -> [DoExprStep] -> ([Deconstruct], [DoExprStep])
partitionMatchingDecons Exp
e [DoExprStep]
r
    partitionMatchingDecons Exp
e (DoExprStep
s : [DoExprStep]
r) = ([DoExprStep] -> [DoExprStep])
-> ([Deconstruct], [DoExprStep]) -> ([Deconstruct], [DoExprStep])
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second (DoExprStep
s DoExprStep -> [DoExprStep] -> [DoExprStep]
forall a. a -> [a] -> [a]
:) (([Deconstruct], [DoExprStep]) -> ([Deconstruct], [DoExprStep]))
-> ([Deconstruct], [DoExprStep]) -> ([Deconstruct], [DoExprStep])
forall a b. (a -> b) -> a -> b
$ Exp -> [DoExprStep] -> ([Deconstruct], [DoExprStep])
partitionMatchingDecons Exp
e [DoExprStep]
r