{- |
Module      :  Camfort.Specification.Units.Annotation
Description :  Annotation with unit information.
Copyright   :  (c) 2017, Dominic Orchard, Andrew Rice, Mistral Contrastin, Matthew Danish
License     :  Apache-2.0

Maintainer  :  dom.orchard@gmail.com
Stability   :  experimental

Defines the 'UnitAnnotation' datatype, which is used for annotating a
'ProgramFile' with units information.
-}

{-# LANGUAGE DeriveDataTypeable    #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Camfort.Specification.Units.Annotation
  (
    -- * Annotation Type
    UA
  , mkUnitAnnotation
  , prevAnnotation
  , unitBlock
  , unitConstraint
  , unitInfo
  , unitPU
  , unitSpec
    -- * Helpers
  , cleanLinks
  , getConstraint
  , getUnitInfo
  , maybeSetUnitConstraintF2
  , maybeSetUnitInfo
  , maybeSetUnitInfoF2
  , setConstraint
  , setUnitInfo
  ) where

import Data.Data (Data, Typeable)
import Data.Generics.Uniplate.Operations (transformBi)

import qualified Language.Fortran.Analysis as FA
import qualified Language.Fortran.AST      as F

import qualified Camfort.Analysis.Annotations             as Ann
import           Camfort.Analysis.CommentAnnotator
  (ASTEmbeddable(..), Linkable(..))
import qualified Camfort.Specification.Units.Environment  as E
import qualified Camfort.Specification.Units.Parser.Types as P

-- The annotation on the AST used for solving units.
data UnitAnnotation a = UnitAnnotation {
    prevAnnotation :: a,
    unitSpec       :: Maybe P.UnitStatement,
    unitConstraint :: Maybe E.Constraint,
    unitInfo       :: Maybe E.UnitInfo,
    unitBlock      :: Maybe (F.Block (FA.Analysis (UnitAnnotation a))), -- ^ linked variable declaration
    unitPU         :: Maybe (F.ProgramUnit (FA.Analysis (UnitAnnotation a))) -- ^ linked program unit
  } deriving (Data, Typeable, Show)

mkUnitAnnotation :: a -> UnitAnnotation a
mkUnitAnnotation a = UnitAnnotation a Nothing Nothing Nothing Nothing Nothing

-- Convenience name for a common annotation type.
type UA = FA.Analysis (UnitAnnotation Ann.A)

-- Instances for embedding parsed specifications into the AST
instance ASTEmbeddable UA P.UnitStatement where
  annotateWithAST ann ast =
    Ann.onPrev (\ann' -> ann' { unitSpec = Just ast }) ann

-- Link annotation comments to declaration statements
instance Linkable UA where
  link ann (b@(F.BlStatement _ _ _ F.StDeclaration {})) =
      Ann.onPrev (\ann' -> ann' { unitBlock = Just b }) ann
  link ann _ = ann
  linkPU ann pu@F.PUFunction{} =
      Ann.onPrev (\ann' -> ann' { unitPU = Just pu }) ann
  linkPU ann pu@F.PUSubroutine{} =
      Ann.onPrev (\ann' -> ann' { unitPU = Just pu }) ann
  linkPU ann _ = ann

-- | Extract the unit info from a given annotated piece of AST.
getUnitInfo :: F.Annotated f => f UA -> Maybe E.UnitInfo
getUnitInfo = unitInfo . FA.prevAnnotation . F.getAnnotation

-- | Extract the constraint from a given annotated piece of AST.
getConstraint :: F.Annotated f => f UA -> Maybe E.Constraint
getConstraint = unitConstraint . FA.prevAnnotation . F.getAnnotation

-- | Set the UnitInfo field on a piece of AST.
setUnitInfo :: F.Annotated f => E.UnitInfo -> f UA -> f UA
setUnitInfo info = F.modifyAnnotation (Ann.onPrev (\ ua -> ua { unitInfo = Just info }))

-- | Set the Constraint field on a piece of AST.
setConstraint :: F.Annotated f => E.Constraint -> f UA -> f UA
setConstraint (E.ConConj []) = id
setConstraint c              =
  F.modifyAnnotation (Ann.onPrev (\ ua -> ua { unitConstraint = Just c }))

--------------------------------------------------

-- Various helper functions for setting the UnitInfo or Constraint of a piece of AST
maybeSetUnitInfo :: F.Annotated f => Maybe E.UnitInfo -> f UA -> f UA
maybeSetUnitInfo u e  = maybe e (`setUnitInfo` e) u

maybeSetUnitInfoF2 :: F.Annotated f => (a -> b -> E.UnitInfo) -> Maybe a -> Maybe b -> f UA -> f UA
maybeSetUnitInfoF2 f u1 u2 e = maybe e (`setUnitInfo` e) (f <$> u1 <*> u2)

maybeSetUnitConstraintF2 :: F.Annotated f => (a -> b -> E.Constraint) -> Maybe a -> Maybe b -> f UA -> f UA
maybeSetUnitConstraintF2 f u1 u2 e = maybe e (`setConstraint` e) (f <$> u1 <*> u2)

cleanLinks :: F.ProgramFile UA -> F.ProgramFile UA
cleanLinks = transformBi
  (\a -> a { unitPU    = Nothing
           , unitBlock = Nothing
           , unitSpec  = Nothing
           } :: UnitAnnotation Ann.A)