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

HIE AST visitor for single-pass traversal.
-}

module Stan.Analysis.Visitor
    ( VisitorState (..)
    , getFinalObservations
    , addObservation
    , addObservations
    , addFixity
    , addOpDecl

    , Visitor (..)
    , visitAst
    ) where

import Relude.Extra.Lens (Lens', lens, over)

import Stan.Ghc.Compat (RealSrcSpan)
import Stan.Hie.Compat (HieAST (..), HieASTs (..), HieFile (..), TypeIndex)
import Stan.Inspection (inspectionId)
import Stan.Inspection.Style (stan0301)
import Stan.Observation (Observation, Observations, mkObservation)

import qualified Data.HashMap.Strict as HM
import qualified Data.Map.Strict as Map
import qualified Slist as S


{- | State for the 'Visitor' object that stores all values during a
single HIE AST traversal.
-}
data VisitorState = VisitorState
    { VisitorState -> Slist Observation
visitorStateObservations :: !Observations

      -- Operators for STAN-0301
    , VisitorState -> HashMap Text ()
visitorStateFixities     :: !(HashMap Text ())
    , VisitorState -> HashMap Text RealSrcSpan
visitorStateOpDecls      :: !(HashMap Text RealSrcSpan)
    }

-- | Initial empty state.
initialVisitorState :: VisitorState
initialVisitorState :: VisitorState
initialVisitorState = VisitorState
    { visitorStateObservations :: Slist Observation
visitorStateObservations = Slist Observation
forall a. Monoid a => a
mempty
    , visitorStateFixities :: HashMap Text ()
visitorStateFixities     = HashMap Text ()
forall a. Monoid a => a
mempty
    , visitorStateOpDecls :: HashMap Text RealSrcSpan
visitorStateOpDecls      = HashMap Text RealSrcSpan
forall a. Monoid a => a
mempty
    }

{- | Transform 'VisitorState' to the final list of observations for
the given 'HieFile'. 'VisitorState' stores not only ready
'Observations' but also additional metadata collected during tree
traversal, so this metadata is converted to 'Observations' for the
corresponding 'Inspection's.
-}
finaliseState :: HieFile -> VisitorState -> Observations
finaliseState :: HieFile -> VisitorState -> Slist Observation
finaliseState HieFile
hie VisitorState{HashMap Text ()
HashMap Text RealSrcSpan
Slist Observation
visitorStateObservations :: VisitorState -> Slist Observation
visitorStateFixities :: VisitorState -> HashMap Text ()
visitorStateOpDecls :: VisitorState -> HashMap Text RealSrcSpan
visitorStateObservations :: Slist Observation
visitorStateFixities :: HashMap Text ()
visitorStateOpDecls :: HashMap Text RealSrcSpan
..} =
    -- STAN-0301: missing fixity declaration
    -- detected by finding a difference between two sets:
    -- 1. Top-level defined operators
    -- 2. Fixity declarations for operators in module
    let operatorsWithoutFixity :: HashMap Text RealSrcSpan
operatorsWithoutFixity = HashMap Text RealSrcSpan
-> HashMap Text () -> HashMap Text RealSrcSpan
forall k v w.
(Eq k, Hashable k) =>
HashMap k v -> HashMap k w -> HashMap k v
HM.difference HashMap Text RealSrcSpan
visitorStateOpDecls HashMap Text ()
visitorStateFixities
        stan0301inss :: Slist Observation
stan0301inss = Id Inspection -> HieFile -> RealSrcSpan -> Observation
mkObservation (Inspection -> Id Inspection
inspectionId Inspection
stan0301) HieFile
hie (RealSrcSpan -> Observation)
-> Slist RealSrcSpan -> Slist Observation
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [RealSrcSpan] -> Slist RealSrcSpan
forall a. [a] -> Slist a
S.slist (HashMap Text RealSrcSpan -> [RealSrcSpan]
forall a. HashMap Text a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList HashMap Text RealSrcSpan
operatorsWithoutFixity)
    -- combine final observations
    in Slist Observation
visitorStateObservations Slist Observation -> Slist Observation -> Slist Observation
forall a. Semigroup a => a -> a -> a
<> Slist Observation
stan0301inss

-- | Get sized list of all 'Observations' from the given HIE file
-- using the created 'Visitor'.
getFinalObservations :: HieFile -> Visitor -> Observations
getFinalObservations :: HieFile -> Visitor -> Slist Observation
getFinalObservations HieFile
hie Visitor
visitor =
    let visitAction :: StateT VisitorState Identity ()
visitAction = (HieAST TypeIndex -> StateT VisitorState Identity ())
-> [HieAST TypeIndex] -> StateT VisitorState Identity ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Visitor -> HieAST TypeIndex -> StateT VisitorState Identity ()
visitAst Visitor
visitor) [HieAST TypeIndex]
allHieAsts
        resultState :: VisitorState
resultState = StateT VisitorState Identity () -> VisitorState -> VisitorState
forall s a. State s a -> s -> s
execState StateT VisitorState Identity ()
visitAction VisitorState
initialVisitorState
    in HieFile -> VisitorState -> Slist Observation
finaliseState HieFile
hie VisitorState
resultState
  where
    allHieAsts :: [HieAST TypeIndex]
    allHieAsts :: [HieAST TypeIndex]
allHieAsts = Map HiePath (HieAST TypeIndex) -> [HieAST TypeIndex]
forall k a. Map k a -> [a]
Map.elems (Map HiePath (HieAST TypeIndex) -> [HieAST TypeIndex])
-> Map HiePath (HieAST TypeIndex) -> [HieAST TypeIndex]
forall a b. (a -> b) -> a -> b
$ HieASTs TypeIndex -> Map HiePath (HieAST TypeIndex)
forall a. HieASTs a -> Map HiePath (HieAST a)
getAsts (HieASTs TypeIndex -> Map HiePath (HieAST TypeIndex))
-> HieASTs TypeIndex -> Map HiePath (HieAST TypeIndex)
forall a b. (a -> b) -> a -> b
$ HieFile -> HieASTs TypeIndex
hie_asts HieFile
hie

observationsL :: Lens' VisitorState Observations
observationsL :: Lens' VisitorState (Slist Observation)
observationsL = (VisitorState -> Slist Observation)
-> (VisitorState -> Slist Observation -> VisitorState)
-> Lens' VisitorState (Slist Observation)
forall s a. (s -> a) -> (s -> a -> s) -> Lens' s a
lens
    VisitorState -> Slist Observation
visitorStateObservations
    (\VisitorState
vstate Slist Observation
new -> VisitorState
vstate { visitorStateObservations = new })

fixitiesL :: Lens' VisitorState (HashMap Text ())
fixitiesL :: Lens' VisitorState (HashMap Text ())
fixitiesL = (VisitorState -> HashMap Text ())
-> (VisitorState -> HashMap Text () -> VisitorState)
-> Lens' VisitorState (HashMap Text ())
forall s a. (s -> a) -> (s -> a -> s) -> Lens' s a
lens
    VisitorState -> HashMap Text ()
visitorStateFixities
    (\VisitorState
vstate HashMap Text ()
new -> VisitorState
vstate { visitorStateFixities = new })

opDeclsL :: Lens' VisitorState (HashMap Text RealSrcSpan)
opDeclsL :: Lens' VisitorState (HashMap Text RealSrcSpan)
opDeclsL = (VisitorState -> HashMap Text RealSrcSpan)
-> (VisitorState -> HashMap Text RealSrcSpan -> VisitorState)
-> Lens' VisitorState (HashMap Text RealSrcSpan)
forall s a. (s -> a) -> (s -> a -> s) -> Lens' s a
lens
    VisitorState -> HashMap Text RealSrcSpan
visitorStateOpDecls
    (\VisitorState
vstate HashMap Text RealSrcSpan
new -> VisitorState
vstate { visitorStateOpDecls = new })

-- | Add single 'Observation' to the existing 'VisitorState'.
addObservation :: Observation -> State VisitorState ()
addObservation :: Observation -> StateT VisitorState Identity ()
addObservation Observation
obs = (VisitorState -> VisitorState) -> StateT VisitorState Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' ((VisitorState -> VisitorState) -> StateT VisitorState Identity ())
-> (VisitorState -> VisitorState)
-> StateT VisitorState Identity ()
forall a b. (a -> b) -> a -> b
$ Lens' VisitorState (Slist Observation)
-> (Slist Observation -> Slist Observation)
-> VisitorState
-> VisitorState
forall s a. Lens' s a -> (a -> a) -> s -> s
over (Slist Observation -> f (Slist Observation))
-> VisitorState -> f VisitorState
Lens' VisitorState (Slist Observation)
observationsL (Observation -> Slist Observation
forall a. a -> Slist a
S.one Observation
obs Slist Observation -> Slist Observation -> Slist Observation
forall a. Semigroup a => a -> a -> a
<>)

-- | Add 'Observations' to the existing 'VisitorState'.
addObservations :: Observations -> State VisitorState ()
addObservations :: Slist Observation -> StateT VisitorState Identity ()
addObservations Slist Observation
obss
    | Slist Observation -> Bool
forall a. Slist a -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Slist Observation
obss = StateT VisitorState Identity ()
forall (f :: * -> *). Applicative f => f ()
pass
    | Bool
otherwise = (VisitorState -> VisitorState) -> StateT VisitorState Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' ((VisitorState -> VisitorState) -> StateT VisitorState Identity ())
-> (VisitorState -> VisitorState)
-> StateT VisitorState Identity ()
forall a b. (a -> b) -> a -> b
$ Lens' VisitorState (Slist Observation)
-> (Slist Observation -> Slist Observation)
-> VisitorState
-> VisitorState
forall s a. Lens' s a -> (a -> a) -> s -> s
over (Slist Observation -> f (Slist Observation))
-> VisitorState -> f VisitorState
Lens' VisitorState (Slist Observation)
observationsL (Slist Observation
obss Slist Observation -> Slist Observation -> Slist Observation
forall a. Semigroup a => a -> a -> a
<>)

-- | Add single operator infix declaration.
addFixity :: Text -> State VisitorState ()
addFixity :: Text -> StateT VisitorState Identity ()
addFixity Text
fixity = (VisitorState -> VisitorState) -> StateT VisitorState Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' ((VisitorState -> VisitorState) -> StateT VisitorState Identity ())
-> (VisitorState -> VisitorState)
-> StateT VisitorState Identity ()
forall a b. (a -> b) -> a -> b
$ Lens' VisitorState (HashMap Text ())
-> (HashMap Text () -> HashMap Text ())
-> VisitorState
-> VisitorState
forall s a. Lens' s a -> (a -> a) -> s -> s
over (HashMap Text () -> f (HashMap Text ()))
-> VisitorState -> f VisitorState
Lens' VisitorState (HashMap Text ())
fixitiesL (Text -> () -> HashMap Text () -> HashMap Text ()
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HM.insert Text
fixity ())

-- | Add single operator top-level defintion with its position.
addOpDecl :: Text -> RealSrcSpan -> State VisitorState ()
addOpDecl :: Text -> RealSrcSpan -> StateT VisitorState Identity ()
addOpDecl Text
opDecl RealSrcSpan
srcSpan = (VisitorState -> VisitorState) -> StateT VisitorState Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' ((VisitorState -> VisitorState) -> StateT VisitorState Identity ())
-> (VisitorState -> VisitorState)
-> StateT VisitorState Identity ()
forall a b. (a -> b) -> a -> b
$ Lens' VisitorState (HashMap Text RealSrcSpan)
-> (HashMap Text RealSrcSpan -> HashMap Text RealSrcSpan)
-> VisitorState
-> VisitorState
forall s a. Lens' s a -> (a -> a) -> s -> s
over (HashMap Text RealSrcSpan -> f (HashMap Text RealSrcSpan))
-> VisitorState -> f VisitorState
Lens' VisitorState (HashMap Text RealSrcSpan)
opDeclsL (Text
-> RealSrcSpan
-> HashMap Text RealSrcSpan
-> HashMap Text RealSrcSpan
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
HM.insert Text
opDecl RealSrcSpan
srcSpan)

-- | Object that implements the /Visitor pattern/.
newtype Visitor = Visitor
    { Visitor -> HieAST TypeIndex -> StateT VisitorState Identity ()
unVisitor :: HieAST TypeIndex -> State VisitorState ()
    }

-- | Traverse HIE AST starting from a given node using 'Visitor'.
visitAst :: Visitor -> HieAST TypeIndex -> State VisitorState ()
visitAst :: Visitor -> HieAST TypeIndex -> StateT VisitorState Identity ()
visitAst (Visitor HieAST TypeIndex -> StateT VisitorState Identity ()
visit) = HieAST TypeIndex -> StateT VisitorState Identity ()
go
  where
    go :: HieAST TypeIndex -> State VisitorState ()
    go :: HieAST TypeIndex -> StateT VisitorState Identity ()
go node :: HieAST TypeIndex
node@Node{[HieAST TypeIndex]
RealSrcSpan
SourcedNodeInfo TypeIndex
sourcedNodeInfo :: SourcedNodeInfo TypeIndex
nodeSpan :: RealSrcSpan
nodeChildren :: [HieAST TypeIndex]
sourcedNodeInfo :: forall a. HieAST a -> SourcedNodeInfo a
nodeSpan :: forall a. HieAST a -> RealSrcSpan
nodeChildren :: forall a. HieAST a -> [HieAST a]
..} = do
        HieAST TypeIndex -> StateT VisitorState Identity ()
visit HieAST TypeIndex
node
        (HieAST TypeIndex -> StateT VisitorState Identity ())
-> [HieAST TypeIndex] -> StateT VisitorState Identity ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ HieAST TypeIndex -> StateT VisitorState Identity ()
go [HieAST TypeIndex]
nodeChildren