{-# LANGUAGE CPP #-}

{- |
Copyright: (c) 2020 Kowainik
SPDX-License-Identifier: MPL-2.0
Maintainer: Kowainik <xrom.xkov@gmail.com>

Patterns for AST and syntax tree nodes search.
-}

module Stan.Pattern.Ast
    ( -- * Type
      PatternAst (..)
    , Literal (..)

      -- * Helpers
    , namesToPatternAst
    , anyNamesToPatternAst

      -- * eDSL
    , app
    , opApp
    , constructor
    , constructorNameIdentifier
    , dataDecl
    , fixity
    , fun
    , guardBranch
    , lazyField
    , range
    , rhs
    , tuple
    , typeSig

      -- * Pattern matching
    , case'
    , lambdaCase
    , patternMatchBranch
    , patternMatchArrow
    , patternMatch_
    , literalPat
    , wildPat

      -- * More low-level interface
    , literalAnns
    ) where

import Stan.Ghc.Compat (FastString)
import Stan.Hie.Compat (DeclType (..))
import Stan.NameMeta (NameMeta (..))
import Stan.Pattern.Edsl (PatternBool (..))
import Stan.Pattern.Type (PatternType)

import qualified Data.Set as Set


{- | Query pattern used to search AST nodes in HIE AST. This data type
tries to mirror HIE AST to each future matching, so it's quite
low-level, but helper functions are provided.
-}
data PatternAst
    -- | Integer constant in code.
    = PatternAstConstant !Literal
    -- | Name of a specific function, variable or data type.
    | PatternAstName !NameMeta !PatternType
    -- | Variable name.
    | PatternAstVarName !String
    -- | AST node with tags for current node and any children.
    | PatternAstNode
        !(Set (FastString, FastString))  -- ^ Set of context info (pairs of tags)
    -- | AST node with tags for current node and children
    -- patterns. This pattern should match the node exactly.
    | PatternAstNodeExact
        !(Set (FastString, FastString))  -- ^ Set of context info (pairs of tags)
        ![PatternAst]  -- ^ Node children
    -- | AST wildcard, matches anything.
    | PatternAstAnything
    -- | Choice between patterns. Should match either of them.
    | PatternAstOr !PatternAst !PatternAst
    -- | Union of patterns. Should match both of them.
    | PatternAstAnd !PatternAst !PatternAst
    -- | Negation of pattern. Should match everything except this pattern.
    | PatternAstNeg !PatternAst
    -- | AST node with the specified Identifier details (only 'DeclType')
    | PatternAstIdentifierDetailsDecl !DeclType
    deriving stock (Int -> PatternAst -> ShowS
[PatternAst] -> ShowS
PatternAst -> String
(Int -> PatternAst -> ShowS)
-> (PatternAst -> String)
-> ([PatternAst] -> ShowS)
-> Show PatternAst
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PatternAst] -> ShowS
$cshowList :: [PatternAst] -> ShowS
show :: PatternAst -> String
$cshow :: PatternAst -> String
showsPrec :: Int -> PatternAst -> ShowS
$cshowsPrec :: Int -> PatternAst -> ShowS
Show, PatternAst -> PatternAst -> Bool
(PatternAst -> PatternAst -> Bool)
-> (PatternAst -> PatternAst -> Bool) -> Eq PatternAst
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: PatternAst -> PatternAst -> Bool
$c/= :: PatternAst -> PatternAst -> Bool
== :: PatternAst -> PatternAst -> Bool
$c== :: PatternAst -> PatternAst -> Bool
Eq)

instance PatternBool PatternAst where
    (?) :: PatternAst
    ? :: PatternAst
(?) = PatternAst
PatternAstAnything

    neg :: PatternAst -> PatternAst
    neg :: PatternAst -> PatternAst
neg = PatternAst -> PatternAst
PatternAstNeg

    (|||) :: PatternAst -> PatternAst -> PatternAst
    ||| :: PatternAst -> PatternAst -> PatternAst
(|||) = PatternAst -> PatternAst -> PatternAst
PatternAstOr

    (&&&) :: PatternAst -> PatternAst -> PatternAst
    &&& :: PatternAst -> PatternAst -> PatternAst
(&&&) = PatternAst -> PatternAst -> PatternAst
PatternAstAnd

data Literal
    = ExactNum !Int
    | ExactStr !ByteString
    | PrefixStr !ByteString
    | ContainStr !ByteString
    | AnyLiteral
    deriving stock (Int -> Literal -> ShowS
[Literal] -> ShowS
Literal -> String
(Int -> Literal -> ShowS)
-> (Literal -> String) -> ([Literal] -> ShowS) -> Show Literal
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Literal] -> ShowS
$cshowList :: [Literal] -> ShowS
show :: Literal -> String
$cshow :: Literal -> String
showsPrec :: Int -> Literal -> ShowS
$cshowsPrec :: Int -> Literal -> ShowS
Show, Literal -> Literal -> Bool
(Literal -> Literal -> Bool)
-> (Literal -> Literal -> Bool) -> Eq Literal
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Literal -> Literal -> Bool
$c/= :: Literal -> Literal -> Bool
== :: Literal -> Literal -> Bool
$c== :: Literal -> Literal -> Bool
Eq)

{- | Function that creates 'PatternAst' from the given non-empty list of pairs
'NameMeta' and 'PatternType'.

If the list contains only one 'PatternType' then it is simple 'PatternAstName'.
Else it is 'PatternAstOr' of all such 'PatternAstName's.
-}
namesToPatternAst :: NonEmpty (NameMeta, PatternType) -> PatternAst
namesToPatternAst :: NonEmpty (NameMeta, PatternType) -> PatternAst
namesToPatternAst ((nm :: NameMeta
nm, pat :: PatternType
pat) :| []) = NameMeta -> PatternType -> PatternAst
PatternAstName NameMeta
nm PatternType
pat
namesToPatternAst ((nm :: NameMeta
nm, pat :: PatternType
pat) :| x :: (NameMeta, PatternType)
x:rest :: [(NameMeta, PatternType)]
rest) = PatternAst -> PatternAst -> PatternAst
PatternAstOr
    (NameMeta -> PatternType -> PatternAst
PatternAstName NameMeta
nm PatternType
pat)
    (NonEmpty (NameMeta, PatternType) -> PatternAst
namesToPatternAst (NonEmpty (NameMeta, PatternType) -> PatternAst)
-> NonEmpty (NameMeta, PatternType) -> PatternAst
forall a b. (a -> b) -> a -> b
$ (NameMeta, PatternType)
x (NameMeta, PatternType)
-> [(NameMeta, PatternType)] -> NonEmpty (NameMeta, PatternType)
forall a. a -> [a] -> NonEmpty a
:| [(NameMeta, PatternType)]
rest)

-- | Like 'namesToPatternAst' but doesn't care about types.
anyNamesToPatternAst :: NonEmpty NameMeta -> PatternAst
anyNamesToPatternAst :: NonEmpty NameMeta -> PatternAst
anyNamesToPatternAst = NonEmpty (NameMeta, PatternType) -> PatternAst
namesToPatternAst (NonEmpty (NameMeta, PatternType) -> PatternAst)
-> (NonEmpty NameMeta -> NonEmpty (NameMeta, PatternType))
-> NonEmpty NameMeta
-> PatternAst
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (NameMeta -> (NameMeta, PatternType))
-> NonEmpty NameMeta -> NonEmpty (NameMeta, PatternType)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (, PatternType
forall a. PatternBool a => a
(?))

-- | @app f x@ is a pattern for function application @f x@.
app :: PatternAst -> PatternAst -> PatternAst
app :: PatternAst -> PatternAst -> PatternAst
app f :: PatternAst
f x :: PatternAst
x = Set (FastString, FastString) -> [PatternAst] -> PatternAst
PatternAstNodeExact (OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("HsApp", "HsExpr")) [PatternAst
f, PatternAst
x]

-- | @opApp x op y@ is a pattern for operator application @x `op` y@.
opApp :: PatternAst -> PatternAst -> PatternAst -> PatternAst
opApp :: PatternAst -> PatternAst -> PatternAst -> PatternAst
opApp x :: PatternAst
x op :: PatternAst
op y :: PatternAst
y = Set (FastString, FastString) -> [PatternAst] -> PatternAst
PatternAstNodeExact (OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("OpApp", "HsExpr")) [PatternAst
x, PatternAst
op, PatternAst
y]

-- | @range a b@ is a pattern for @[a .. b]@
range :: PatternAst -> PatternAst -> PatternAst
range :: PatternAst -> PatternAst -> PatternAst
range from :: PatternAst
from to :: PatternAst
to = Set (FastString, FastString) -> [PatternAst] -> PatternAst
PatternAstNodeExact (OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("ArithSeq", "HsExpr")) [PatternAst
from, PatternAst
to]

-- | 'lambdaCase' is a pattern for @\case@ expression (not considering branches).
lambdaCase :: PatternAst
lambdaCase :: PatternAst
lambdaCase = Set (FastString, FastString) -> PatternAst
PatternAstNode (OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("HsLamCase", "HsExpr"))

-- | 'case'' is a pattern for @case EXP of@ expression (not considering branches).
case' :: PatternAst
case' :: PatternAst
case' = Set (FastString, FastString) -> PatternAst
PatternAstNode (OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("HsCase", "HsExpr"))

-- | Pattern to represent one pattern matching branch.
patternMatchBranch :: PatternAst
patternMatchBranch :: PatternAst
patternMatchBranch = Set (FastString, FastString) -> PatternAst
PatternAstNode (OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("Match", "Match"))

{- | Pattern for @_@ in pattern matching.

__Note:__ presents on GHC >=8.10 only.
-}
wildPat :: PatternAst
wildPat :: PatternAst
wildPat = Set (FastString, FastString) -> PatternAst
PatternAstNode (OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("WildPat", "Pat"))

{- | Pattern for literals in pattern matching.

__Note:__ presents on GHC >=8.10 only.
-}
literalPat :: PatternAst
literalPat :: PatternAst
literalPat = Set (FastString, FastString) -> PatternAst
PatternAstNode (OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("NPat", "Pat"))
    PatternAst -> PatternAst -> PatternAst
forall a. PatternBool a => a -> a -> a
||| Set (FastString, FastString) -> PatternAst
PatternAstNode (OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("LitPat", "Pat"))

-- | Pattern to represent one pattern matching branch on @_@.
patternMatch_ :: PatternAst -> PatternAst
patternMatch_ :: PatternAst -> PatternAst
patternMatch_ val :: PatternAst
val = Set (FastString, FastString) -> [PatternAst] -> PatternAst
PatternAstNodeExact (OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("Match", "Match"))
#if __GLASGOW_HASKELL__ >= 810
    $ wildPat :
#endif
    [PatternAst -> PatternAst
patternMatchArrow PatternAst
val]

-- | Pattern to represent right side of the pattern matching, e.g. @-> "foo"@.
patternMatchArrow :: PatternAst -> PatternAst
patternMatchArrow :: PatternAst -> PatternAst
patternMatchArrow x :: PatternAst
x = Set (FastString, FastString) -> [PatternAst] -> PatternAst
PatternAstNodeExact (OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("GRHS", "GRHS")) [PatternAst
x]

{- | Pattern for the top-level fixity declaration:

@
infixr 7 ***, +++, ???
@
-}
fixity :: PatternAst
fixity :: PatternAst
fixity = Set (FastString, FastString) -> PatternAst
PatternAstNode (Set (FastString, FastString) -> PatternAst)
-> Set (FastString, FastString) -> PatternAst
forall a b. (a -> b) -> a -> b
$ OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("FixitySig", "FixitySig")

{- | Pattern for the function type signature declaration:

@
foo :: Some -> Type
@
-}
typeSig :: PatternAst
typeSig :: PatternAst
typeSig = Set (FastString, FastString) -> PatternAst
PatternAstNode (Set (FastString, FastString) -> PatternAst)
-> Set (FastString, FastString) -> PatternAst
forall a b. (a -> b) -> a -> b
$ OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("TypeSig", "Sig")

{- | Pattern for the function definition:

@
foo x y = ...
@
-}
fun :: PatternAst
fun :: PatternAst
fun = Set (FastString, FastString) -> PatternAst
PatternAstNode (Set (FastString, FastString) -> PatternAst)
-> Set (FastString, FastString) -> PatternAst
forall a b. (a -> b) -> a -> b
$ [(FastString, FastString)] -> Set (FastString, FastString)
forall a. Ord a => [a] -> Set a
Set.fromList
    [ ("AbsBinds", "HsBindLR")
    , ("FunBind",  "HsBindLR")
    , ("Match",    "Match")
    ]

{- | @data@ or @newtype@ declaration.
-}
dataDecl :: PatternAst
dataDecl :: PatternAst
dataDecl = Set (FastString, FastString) -> PatternAst
PatternAstNode (Set (FastString, FastString) -> PatternAst)
-> Set (FastString, FastString) -> PatternAst
forall a b. (a -> b) -> a -> b
$ OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("DataDecl", "TyClDecl")

{- | Constructor of a plain data type or newtype. Children of node
that matches this pattern are constructor fields.
-}
constructor :: PatternAst
constructor :: PatternAst
constructor = Set (FastString, FastString) -> PatternAst
PatternAstNode (Set (FastString, FastString) -> PatternAst)
-> Set (FastString, FastString) -> PatternAst
forall a b. (a -> b) -> a -> b
$ OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("ConDeclH98", "ConDecl")

{- | Constructor name Identifier info
-}
constructorNameIdentifier :: PatternAst
constructorNameIdentifier :: PatternAst
constructorNameIdentifier = DeclType -> PatternAst
PatternAstIdentifierDetailsDecl DeclType
ConDec

{- | Lazy data type field. Comes in two shapes:

1. Record field, like: @foo :: Text@
2. Simple type: @Int@
-}
lazyField :: PatternAst
lazyField :: PatternAst
lazyField = PatternAst
lazyRecordField PatternAst -> PatternAst -> PatternAst
forall a. PatternBool a => a -> a -> a
||| PatternAst
type_

{- | Pattern for any occurrence of a plain type. Covers the following
cases:

* Simple type: Int, Bool, a
* Higher-kinded type: Maybe Int, Either String a
* Type in parenthesis: (Int)
* Tuples: (Int, Bool)
* List type: [Int]
* Function type: Int -> Bool
-}
type_ :: PatternAst
type_ :: PatternAst
type_ =
    Set (FastString, FastString) -> PatternAst
PatternAstNode (OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("HsTyVar", "HsType"))  -- simple type: Int, Bool
    PatternAst -> PatternAst -> PatternAst
forall a. PatternBool a => a -> a -> a
|||
    Set (FastString, FastString) -> PatternAst
PatternAstNode (OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("HsAppTy", "HsType"))  -- composite: Maybe Int
    PatternAst -> PatternAst -> PatternAst
forall a. PatternBool a => a -> a -> a
|||
    Set (FastString, FastString) -> PatternAst
PatternAstNode (OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("HsParTy", "HsType"))  -- type in ()
    PatternAst -> PatternAst -> PatternAst
forall a. PatternBool a => a -> a -> a
|||
    Set (FastString, FastString) -> PatternAst
PatternAstNode (OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("HsTupleTy", "HsType"))  -- tuple types: (Int, Bool)
    PatternAst -> PatternAst -> PatternAst
forall a. PatternBool a => a -> a -> a
|||
    Set (FastString, FastString) -> PatternAst
PatternAstNode (OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("HsListTy", "HsType"))  -- list types: [Int]
    PatternAst -> PatternAst -> PatternAst
forall a. PatternBool a => a -> a -> a
|||
    Set (FastString, FastString) -> PatternAst
PatternAstNode (OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("HsFunTy", "HsType"))  -- function types: Int -> Bool

{- | Pattern for the field without the explicit bang pattern:

@
someField :: Int
@
-}
lazyRecordField :: PatternAst
lazyRecordField :: PatternAst
lazyRecordField = Set (FastString, FastString) -> [PatternAst] -> PatternAst
PatternAstNodeExact
    (OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("ConDeclField", "ConDeclField"))
    [ Set (FastString, FastString) -> PatternAst
PatternAstNode
        ([Item (Set (FastString, FastString))]
-> Set (FastString, FastString)
forall l. IsList l => [Item l] -> l
fromList
            [ ("AbsBinds", "HsBindLR")
            , ("FunBind", "HsBindLR")
            ]
        )
    , PatternAst
type_
    ]

{- | Pattern for tuples:

* Type signatures: foo :: (Int, Int, Int, Int)
* Literals: (True, 0, [], Nothing)
-}
tuple :: PatternAst
tuple :: PatternAst
tuple =
    Set (FastString, FastString) -> PatternAst
PatternAstNode (OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("HsTupleTy", "HsType"))  -- tuple type
    PatternAst -> PatternAst -> PatternAst
forall a. PatternBool a => a -> a -> a
|||
    Set (FastString, FastString) -> PatternAst
PatternAstNode (OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("ExplicitTuple", "HsExpr"))  -- tuple literal

{- | Pattern for a single @guard@ branch:

@
    | x < y = ...
@
-}
guardBranch :: PatternAst
guardBranch :: PatternAst
guardBranch = Set (FastString, FastString) -> PatternAst
PatternAstNode (Set (FastString, FastString) -> PatternAst)
-> Set (FastString, FastString) -> PatternAst
forall a b. (a -> b) -> a -> b
$ OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("BodyStmt", "StmtLR")

{- | Pattern for the right-hand-side. Usually an equality sign.

@
   foo = baz
@
-}
rhs :: PatternAst
rhs :: PatternAst
rhs = Set (FastString, FastString) -> PatternAst
PatternAstNode (Set (FastString, FastString) -> PatternAst)
-> Set (FastString, FastString) -> PatternAst
forall a b. (a -> b) -> a -> b
$ OneItem (Set (FastString, FastString))
-> Set (FastString, FastString)
forall x. One x => OneItem x -> x
one ("GRHS", "GRHS")

-- | Annotations for constants: 0, "foo", etc.
literalAnns :: (FastString, FastString)
literalAnns :: (FastString, FastString)
literalAnns = ("HsOverLit", "HsExpr")