content: {-
     Copyright 2016, Dominic Orchard, Andrew Rice, Mistral Contrastin, Matthew Danish
  
     Licensed under the Apache License, Version 2.0 (the "License");
     you may not use this file except in compliance with the License.
     You may obtain a copy of the License at
  
         http://www.apache.org/licenses/LICENSE-2.0
  
     Unless required by applicable law or agreed to in writing, software
     distributed under the License is distributed on an "AS IS" BASIS,
     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     See the License for the specific language governing permissions and
     limitations under the License.
  -}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE OverloadedStrings #-}
  
  module Camfort.Transformation.EquivalenceElim
    ( refactorEquivalences
    ) where
  
  import           Camfort.Analysis
  import           Camfort.Analysis.Annotations
  import           Camfort.Helpers.Syntax
  import           Camfort.Transformation.DeadCode
  import           Control.Monad.State.Lazy hiding (ap)
  import           Data.Generics.Uniplate.Operations
  import           Data.List
  import qualified Data.Map as M
  import           Data.Void (Void)
  import qualified Language.Fortran.AST as F
  import qualified Language.Fortran.Analysis as FA
  import qualified Language.Fortran.Analysis.Renaming as FAR
  import qualified Language.Fortran.Analysis.Types as FAT (analyseTypes, TypeEnv)
  import qualified Language.Fortran.Util.Position as FU
  
  import qualified Debug.Trace
  
  type EquivalenceRefactoring = PureAnalysis Void Void
  
  type A1 = FA.Analysis Annotation
  type RmEqState = ([[F.Expression A1]], Int)
  
  refactorEquivalences :: F.ProgramFile A -> EquivalenceRefactoring (F.ProgramFile A)
  refactorEquivalences pf = do
    let
      -- initialise analysis
      pf'             = FAR.analyseRenames . FA.initAnalysis $ pf
      -- calculate types
      (pf'', typeEnv) = FAT.analyseTypes pf'
      -- Remove equivalences and add appropriate copy statements
    pf''' <- refactoring typeEnv pf''
    -- Lastly deadcode eliminate any redundant copy statements
    -- generated by the refactoring (but don't dead code elim
    -- existing code)
    deadCode True (fmap FA.prevAnnotation pf''')
    where
      refactoring
        :: FAT.TypeEnv -> F.ProgramFile A1
        -> EquivalenceRefactoring (F.ProgramFile A1)
      refactoring tenv pf' = do
        (pf'', _) <- runStateT equiv ([], 0)
        return pf''
        where
           equiv = do pf'' <- transformBiM perBlockRmEquiv pf'
                      descendBiM (addCopysPerBlockGroup tenv) pf''
  
  addCopysPerBlockGroup
    :: FAT.TypeEnv -> [F.Block A1]
    -> StateT RmEqState EquivalenceRefactoring [F.Block A1]
  addCopysPerBlockGroup tenv blocks = do
      blockss <- mapM (addCopysPerBlock tenv) blocks
      return $ concat blockss
  
  addCopysPerBlock
    :: FAT.TypeEnv -> F.Block A1
    -> StateT RmEqState EquivalenceRefactoring [F.Block A1]
  addCopysPerBlock tenv b@(F.BlStatement _ _ _
                   (F.StExpressionAssign a sp@(FU.SrcSpan s1 _) dstE _))
    | not (pRefactored $ FA.prevAnnotation a) = do
      -- Find all variables/cells that are equivalent to the target
      -- of this assignment
      eqs <- equivalentsToExpr dstE
      Debug.Trace.trace (show (length eqs)) $
          -- If there is only one, then it must refer to itself, so do nothing
          if length eqs <= 1
            then return [b]
          -- If there are more than one, copy statements must be generated
            else do
              (equivs, n) <- get
  
              -- Remove the destination from the equivalents
              let eqs' = deleteBy (\ x y -> af x == af y) dstE eqs
  
              -- Make copy statements
              let pos = afterAligned sp
              let copies = map (mkCopy tenv pos dstE) eqs'
  
              let (FU.Position ao c l f p) = s1
                  reportSpan i =
                    let pos' = FU.Position (ao + i) c (l + i) f p
                    in (FU.SrcSpan pos' pos')
  
              forM_ [n..(n + length copies - 1)] $ \i -> do
                origin <- atSpanned (reportSpan i)
                logInfo origin $ "added copy due to refactored equivalence"
  
              -- Update refactoring state
              put (equivs, n + length eqs')
              -- Sequence original assignment with new assignments
              return $ b : copies
  
  addCopysPerBlock tenv x = do
     x' <- descendBiM (addCopysPerBlockGroup tenv) x
     return [x']
  
  -- see if two expressions have the same type
  equalTypes :: FAT.TypeEnv -> F.Expression A1 -> F.Expression A1 -> Maybe FA.IDType
  equalTypes tenv e e' = do
      v1 <- extractVariable e
      v2 <- extractVariable e'
      t1 <- M.lookup v1 tenv
      t2 <- M.lookup v2 tenv
      if t1 == t2 then Just t1 else Nothing
  
  -- Create copy statements. Parameters:
  --    * A type environment to find out if a type cast is needed
  --    * A SrcPos where the copy statements are going to inserted at
  --    * The source expression
  --    * The number of copies to increment the line by
  --           paired with the destination expression
  mkCopy :: FAT.TypeEnv
         -> FU.Position
         -> F.Expression A1 -> F.Expression A1 -> F.Block A1
  mkCopy tenv pos srcE dstE = FA.initAnalysis $
     F.BlStatement a sp Nothing $
       case equalTypes tenv srcE dstE of
         -- Types not equal, so create a transfer
         Nothing -> F.StExpressionAssign a sp dstE' call
                      where
                       call = F.ExpFunctionCall a sp transf argst
                       transf = F.ExpValue a sp (F.ValVariable "transfer")
                       argst  = F.AList a sp args
                       args   = map (F.Argument a sp Nothing . F.ArgExpr) [srcE', dstE']
         -- Types are equal, simple a assignment
         Just _ -> F.StExpressionAssign a sp dstE' srcE'
    where
       -- Set position to be at col = 0
       sp   = FU.SrcSpan (toCol0 pos) (toCol0 pos)
       -- But store the aligned position in refactored so
       -- that the reprint algorithm can add the appropriate indentation
       a = unitAnnotation { refactored = Just pos, newNode = True }
       dstE' = FA.stripAnalysis dstE
       srcE' = FA.stripAnalysis srcE
  
  perBlockRmEquiv :: F.Block A1 -> StateT RmEqState EquivalenceRefactoring (F.Block A1)
  perBlockRmEquiv = transformBiM perStatementRmEquiv
  
  perStatementRmEquiv
    :: F.Statement A1
    -> StateT RmEqState EquivalenceRefactoring (F.Statement A1)
  perStatementRmEquiv (F.StEquivalence a sp@(FU.SrcSpan spL _) equivs) = do
      (ess, n) <- get
  
      let spL' = FU.SrcSpan spL spL
      logInfo' spL' $ "removed equivalence"
  
      put (((map F.aStrip) . F.aStrip $ equivs) ++ ess, n - 1)
      let a' = onPrev (\ap -> ap {refactored = Just spL, deleteNode = True}) a
      return (F.StEquivalence a' (deleteLine sp) equivs)
  perStatementRmEquiv f = return f
  
  -- 'equivalents e' returns a list of variables/memory cells
  -- that have been equivalenced with "e".
  equivalentsToExpr
    :: F.Expression A1
    -> StateT RmEqState EquivalenceRefactoring [F.Expression A1]
  equivalentsToExpr y = do
      (equivs, _) <- get
      return (inGroup y equivs)
    where
      inGroup _ [] = []
      inGroup x (xs:xss) =
          if AnnotationFree x `elem` map AnnotationFree xs
          then xs
          else inGroup x xss
  