{-# LANGUAGE QualifiedDo #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wno-orphans #-}

module Language.Haskell.TH.Natural.Syntax.Expr.Untyped.Class (
    IsExprBuilder (..),

    -- * Let
    letBind,
    strictLetBind,
    letBind_,

    -- * Deconstruction
    getField,
    getTupleField,
    getFields,
    getTupleFields,
    getField',
    getField_,
    getTupleField',
    getField'',
    strict,
) where

import Control.Monad
import qualified Language.Haskell.TH as TH
import Language.Haskell.TH.Gen
import Language.Haskell.TH.Natural.Internal.Utils
import Language.Haskell.TH.Natural.Syntax.Builder hiding (fail)
import Language.Haskell.TH.Natural.Syntax.Expr.Internal

-- | Typeclass to factorise the common behaviour of expression builders
class IsExprBuilder st where
    type Definition st

    addDeconstruct :: Deconstruct -> Builder st step Empty ()

    addLet :: Binding -> Builder st step Empty ()
    letCount :: Builder st step step Int

    returns :: (GenExpr b) => b -> Builder st step Ready ()
    runExprBuilder :: Builder st step Ready () -> Definition st

instance (IsExprBuilder st, GenExpr (Definition st)) => GenExpr (Builder st step Ready ()) where
    genExpr :: Builder st step 'Ready () -> Q Exp
genExpr = Definition st -> Q Exp
forall a. GenExpr a => a -> Q Exp
genExpr (Definition st -> Q Exp)
-> (Builder st step 'Ready () -> Definition st)
-> Builder st step 'Ready ()
-> Q Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder st step 'Ready () -> Definition st
forall st (step :: BuilderStep).
IsExprBuilder st =>
Builder st step 'Ready () -> Definition st
forall (step :: BuilderStep).
Builder st step 'Ready () -> Definition st
runExprBuilder

-- | Let-Bind an expression
letBind :: (IsExprBuilder st, GenExpr b) => b -> Builder st step Empty TH.Exp
letBind :: forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
b -> Builder st step 'Empty Exp
letBind = Bool -> b -> Builder st step 'Empty Exp
forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
Bool -> b -> Builder st step 'Empty Exp
letBind_ Bool
False

-- | Let-Bind an expression with a strict pattern
strictLetBind :: (IsExprBuilder st, GenExpr b) => b -> Builder st step Empty TH.Exp
strictLetBind :: forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
b -> Builder st step 'Empty Exp
strictLetBind = Bool -> b -> Builder st step 'Empty Exp
forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
Bool -> b -> Builder st step 'Empty Exp
letBind_ Bool
True

letBind_ :: (IsExprBuilder st, GenExpr b) => Bool -> b -> Builder st step Empty TH.Exp
letBind_ :: forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
Bool -> b -> Builder st step 'Empty Exp
letBind_ Bool
isStrict b
b = BaseBuilder Q st 'Empty 'Empty Exp
-> BaseBuilder Q st step 'Empty Exp
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 st 'Empty 'Empty Exp
 -> BaseBuilder Q st step 'Empty Exp)
-> BaseBuilder Q st 'Empty 'Empty Exp
-> BaseBuilder Q st step 'Empty Exp
forall a b. (a -> b) -> a -> b
$ do
    prevLetCount <- Builder st 'Empty 'Empty Int
forall st {k} (step :: k).
IsExprBuilder st =>
Builder st step step Int
forall {k} (step :: k). Builder st step step Int
letCount
    bindName <- liftB $ TH.newName ("var" ++ show prevLetCount)
    expr <- liftB $ genExpr b
    addLet $ MkBind bindName expr isStrict
    return $ TH.VarE bindName

getField'' ::
    (IsExprBuilder st, GenExpr b) =>
    -- | The constructor used to deconstruct
    Either Int TH.Name ->
    -- | The index of the field in the constructor
    Int ->
    -- | The number of fields in the constructor, if known
    -- | The expression to deconstruct
    Maybe Int ->
    b ->
    -- | Modify the created pattern (e.g. add BangP or type annotation)
    (TH.Pat -> TH.Q TH.Pat) ->
    Builder st step Empty TH.Exp
getField'' :: forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
Either Int Name
-> Int
-> Maybe Int
-> b
-> (Pat -> Q Pat)
-> Builder st step 'Empty Exp
getField'' Either Int Name
cName Int
idx Maybe Int
fCount b
qExpr Pat -> Q Pat
fPat = BaseBuilder Q st 'Empty 'Empty Exp
-> BaseBuilder Q st step 'Empty Exp
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 st 'Empty 'Empty Exp
 -> BaseBuilder Q st step 'Empty Exp)
-> BaseBuilder Q st 'Empty 'Empty Exp
-> BaseBuilder Q st step 'Empty Exp
forall a b. (a -> b) -> a -> b
$ do
    expr <- Q Exp -> BaseBuilder Q st 'Empty 'Empty Exp
forall {k} (m :: * -> *) a s (step :: k).
Monad m =>
m a -> BaseBuilder m s step step a
liftB (Q Exp -> BaseBuilder Q st 'Empty 'Empty Exp)
-> Q Exp -> BaseBuilder Q st 'Empty 'Empty Exp
forall a b. (a -> b) -> a -> b
$ b -> Q Exp
forall a. GenExpr a => a -> Q Exp
genExpr b
qExpr
    patVarName <- liftB $ TH.newName "pat"
    pat <- liftB $ genPat $ fPat $ TH.VarP patVarName
    fieldCount <- liftB $ maybe (either pure conFieldCount cName) pure fCount
    addDeconstruct $ MkDec cName [(idx, pat)] expr fieldCount
    return $ TH.VarE patVarName

-- | Deconstruct a value and get the nth field of the constructor
getField ::
    (IsExprBuilder st, GenExpr b) =>
    -- | Constructor name
    TH.Name ->
    -- | index of the field
    Int ->
    -- | Expr to deconstruct
    b ->
    Builder st step Empty TH.Exp
getField :: forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
Name -> Int -> b -> Builder st step 'Empty Exp
getField Name
cName Int
idx b
qExpr = Either Int Name
-> Int
-> Maybe Int
-> b
-> (Pat -> Q Pat)
-> Builder st step 'Empty Exp
forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
Either Int Name
-> Int
-> Maybe Int
-> b
-> (Pat -> Q Pat)
-> Builder st step 'Empty Exp
getField'' (Name -> Either Int Name
forall a b. b -> Either a b
Right Name
cName) Int
idx Maybe Int
forall a. Maybe a
Nothing b
qExpr Pat -> Q Pat
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

-- | Same as 'getField', but allow customising the bound pattern
getField' :: (IsExprBuilder st, GenExpr b) => TH.Name -> Int -> b -> (TH.Pat -> TH.Q TH.Pat) -> Builder st step Empty TH.Exp
getField' :: forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
Name -> Int -> b -> (Pat -> Q Pat) -> Builder st step 'Empty Exp
getField' Name
cName Int
idx = Either Int Name
-> Int
-> Maybe Int
-> b
-> (Pat -> Q Pat)
-> Builder st step 'Empty Exp
forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
Either Int Name
-> Int
-> Maybe Int
-> b
-> (Pat -> Q Pat)
-> Builder st step 'Empty Exp
getField'' (Name -> Either Int Name
forall a b. b -> Either a b
Right Name
cName) Int
idx Maybe Int
forall a. Maybe a
Nothing

-- | Similar to 'getField'. Useful when the constructor to use isn't accessible through 'reify'
getField_ ::
    (IsExprBuilder st, GenExpr b) =>
    TH.Name ->
    -- | The index of the field to get
    Int ->
    -- | The number of fields in the constructor
    Int ->
    b ->
    Builder st step Empty TH.Exp
getField_ :: forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
Name -> Int -> Int -> b -> Builder st step 'Empty Exp
getField_ Name
cName Int
idx Int
fCount b
b = Either Int Name
-> Int
-> Maybe Int
-> b
-> (Pat -> Q Pat)
-> Builder st step 'Empty Exp
forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
Either Int Name
-> Int
-> Maybe Int
-> b
-> (Pat -> Q Pat)
-> Builder st step 'Empty Exp
getField'' (Name -> Either Int Name
forall a b. b -> Either a b
Right Name
cName) Int
idx (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
fCount) b
b Pat -> Q Pat
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

-- | Like 'getField', but for tuples
getTupleField ::
    (IsExprBuilder st, GenExpr b) =>
    -- | Tuple size
    Int ->
    -- | Field Index
    Int ->
    b ->
    Builder st step Empty TH.Exp
getTupleField :: forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
Int -> Int -> b -> Builder st step 'Empty Exp
getTupleField Int
size Int
idx b
qExpr = Either Int Name
-> Int
-> Maybe Int
-> b
-> (Pat -> Q Pat)
-> Builder st step 'Empty Exp
forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
Either Int Name
-> Int
-> Maybe Int
-> b
-> (Pat -> Q Pat)
-> Builder st step 'Empty Exp
getField'' (Int -> Either Int Name
forall a b. a -> Either a b
Left Int
size) Int
idx (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
size) b
qExpr Pat -> Q Pat
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

-- | Like 'getField\'', but for tuples
getTupleField' ::
    (IsExprBuilder st, GenExpr b) =>
    -- | Tuple size
    Int ->
    -- | Field Index
    Int ->
    b ->
    (TH.Pat -> TH.Q TH.Pat) ->
    Builder st step Empty TH.Exp
getTupleField' :: forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
Int -> Int -> b -> (Pat -> Q Pat) -> Builder st step 'Empty Exp
getTupleField' Int
size Int
idx = Either Int Name
-> Int
-> Maybe Int
-> b
-> (Pat -> Q Pat)
-> Builder st step 'Empty Exp
forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
Either Int Name
-> Int
-> Maybe Int
-> b
-> (Pat -> Q Pat)
-> Builder st step 'Empty Exp
getField'' (Int -> Either Int Name
forall a b. a -> Either a b
Left Int
size) Int
idx Maybe Int
forall a. Maybe a
Nothing

-- | Deconstruct and get all the fields in the constructor
getFields ::
    (IsExprBuilder st, GenExpr b) =>
    TH.Name ->
    -- | Field count
    Int ->
    b ->
    Builder st step Empty [TH.Exp]
getFields :: forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
Name -> Int -> b -> Builder st step 'Empty [Exp]
getFields Name
n Int
fcount b
b = BaseBuilder Q st 'Empty 'Empty [Exp]
-> BaseBuilder Q st step 'Empty [Exp]
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 st 'Empty 'Empty [Exp]
 -> BaseBuilder Q st step 'Empty [Exp])
-> BaseBuilder Q st 'Empty 'Empty [Exp]
-> BaseBuilder Q st step 'Empty [Exp]
forall a b. (a -> b) -> a -> b
$ do
    e <- Q Exp -> BaseBuilder Q st 'Empty 'Empty Exp
forall {k} (m :: * -> *) a s (step :: k).
Monad m =>
m a -> BaseBuilder m s step step a
liftB (Q Exp -> BaseBuilder Q st 'Empty 'Empty Exp)
-> Q Exp -> BaseBuilder Q st 'Empty 'Empty Exp
forall a b. (a -> b) -> a -> b
$ b -> Q Exp
forall a. GenExpr a => a -> Q Exp
genExpr b
b
    forM [0 .. fcount - 1] $ \Int
i -> Name -> Int -> Int -> Exp -> BaseBuilder Q st 'Empty 'Empty Exp
forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
Name -> Int -> Int -> b -> Builder st step 'Empty Exp
getField_ Name
n Int
i Int
fcount Exp
e

-- | Same as 'getFields', but for tuples
getTupleFields ::
    (IsExprBuilder st, GenExpr b) =>
    -- | Tuple size
    Int ->
    b ->
    Builder st step Empty [TH.Exp]
getTupleFields :: forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
Int -> b -> Builder st step 'Empty [Exp]
getTupleFields Int
size b
b = BaseBuilder Q st 'Empty 'Empty [Exp]
-> BaseBuilder Q st step 'Empty [Exp]
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 st 'Empty 'Empty [Exp]
 -> BaseBuilder Q st step 'Empty [Exp])
-> BaseBuilder Q st 'Empty 'Empty [Exp]
-> BaseBuilder Q st step 'Empty [Exp]
forall a b. (a -> b) -> a -> b
$ do
    e <- Q Exp -> BaseBuilder Q st 'Empty 'Empty Exp
forall {k} (m :: * -> *) a s (step :: k).
Monad m =>
m a -> BaseBuilder m s step step a
liftB (Q Exp -> BaseBuilder Q st 'Empty 'Empty Exp)
-> Q Exp -> BaseBuilder Q st 'Empty 'Empty Exp
forall a b. (a -> b) -> a -> b
$ b -> Q Exp
forall a. GenExpr a => a -> Q Exp
genExpr b
b
    forM [0 .. size - 1] $ \Int
i -> Int -> Int -> Exp -> BaseBuilder Q st 'Empty 'Empty Exp
forall st b (step :: BuilderStep).
(IsExprBuilder st, GenExpr b) =>
Int -> Int -> b -> Builder st step 'Empty Exp
getTupleField Int
size Int
i Exp
e

-- | Util for 'getField\'' to set the pattern as strict
strict :: TH.Pat -> TH.Q TH.Pat
strict :: Pat -> Q Pat
strict = Pat -> Q Pat
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Pat -> Q Pat) -> (Pat -> Pat) -> Pat -> Q Pat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Pat -> Pat
TH.BangP