{-
 - Copyright (c) 2009-2010, ERICSSON AB All rights reserved.
 - 
 - Redistribution and use in source and binary forms, with or without
 - modification, are permitted provided that the following conditions
 - are met:
 - 
 -     * Redistributions of source code must retain the above copyright
 -     notice,
 -       this list of conditions and the following disclaimer.
 -     * Redistributions in binary form must reproduce the above copyright
 -       notice, this list of conditions and the following disclaimer
 -       in the documentation and/or other materials provided with the
 -       distribution.
 -     * Neither the name of the ERICSSON AB nor the names of its
 -     contributors
 -       may be used to endorse or promote products derived from this
 -       software without specific prior written permission.
 - 
 - THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 - "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 - LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
 - A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
 - HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
 - SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 - LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 - DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 - THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 - (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 - OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 -}

{-# LANGUAGE EmptyDataDecls, TypeFamilies, FlexibleInstances #-}

module Feldspar.Compiler.Plugins.ForwardPropagation where

import Feldspar.Compiler.PluginArchitecture
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.List as List
import Feldspar.Compiler.Plugins.PropagationUtils
import Feldspar.Compiler.Error
import Feldspar.Compiler.Options
import Feldspar.Compiler.Imperative.CodeGeneration (simpleType)

fwdPropError = handleError "PluginArch/ForwardPropagation" InternalError

-- ===========================================================================
-- == Copy propagation plugin (forward)
-- ===========================================================================

type VarStatFwd = VarStatistics (Expression ForwardPropagationpSemInf,[VariableData],Bool)
type OccurrencesFwd = Occurrences (Expression ForwardPropagationpSemInf,[VariableData],Bool)

data ForwardPropagation = ForwardPropagation

instance Plugin ForwardPropagation where
    type ExternalInfo ForwardPropagation = DebugOption
    executePlugin ForwardPropagation externalInfo procedure 
		| externalInfo == NoSimplification || externalInfo == NoPrimitiveInstructionHandling = procedure
        | otherwise = fst $ executeTransformationPhase ForwardPropagationTransform (fst globals1) procedureCollected1
            where 
                (procedureCollected1,globals1) = executeTransformationPhase ForwardPropagationCollect Occurrence_read procedure

instance TransformationPhase ForwardPropagation where
    type From ForwardPropagation = ()
    type To ForwardPropagation = ()
    type Downwards ForwardPropagation = ()
    type Upwards ForwardPropagation = ()

-- ====================
--       Collect
-- ====================

data ForwardPropagationpSemInf

instance SemanticInfo ForwardPropagationpSemInf where
    type ProcedureInfo ForwardPropagationpSemInf = ()
    type BlockInfo ForwardPropagationpSemInf = VarStatFwd
    type ProgramInfo ForwardPropagationpSemInf = ()
    type EmptyInfo ForwardPropagationpSemInf = ()
    type PrimitiveInfo ForwardPropagationpSemInf = ()
    type SequenceInfo ForwardPropagationpSemInf = ()
    type BranchInfo ForwardPropagationpSemInf = ()
    type SequentialLoopInfo ForwardPropagationpSemInf = VarStatFwd
    type ParallelLoopInfo ForwardPropagationpSemInf = ()
    type FormalParameterInfo ForwardPropagationpSemInf = ()
    type LocalDeclarationInfo ForwardPropagationpSemInf = ()
    type LeftValueExpressionInfo ForwardPropagationpSemInf = ()
    type VariableInLeftValueInfo ForwardPropagationpSemInf = ()
    type ArrayElemReferenceInfo ForwardPropagationpSemInf = Maybe VariableData --name of the indexed variable
    type InputActualParameterInfo ForwardPropagationpSemInf = ()
    type OutputActualParameterInfo ForwardPropagationpSemInf = ()
    type AssignmentInfo ForwardPropagationpSemInf = ()
    type ProcedureCallInfo ForwardPropagationpSemInf = ()
    type FunctionCallInfo ForwardPropagationpSemInf = ()
    type IntConstantInfo ForwardPropagationpSemInf = ()
    type FloatConstantInfo ForwardPropagationpSemInf = ()
    type BoolConstantInfo ForwardPropagationpSemInf = ()
    type ArrayConstantInfo ForwardPropagationpSemInf = ()
    type VariableInfo ForwardPropagationpSemInf = Occurrence_place

instance Default (Maybe VariableData) where
    defaultValue = Nothing

instance Combine (VarStatFwd, Maybe VariableData) where
    combine a b = (combine (fst a) $ fst b, Nothing)

data ForwardPropagationCollect = ForwardPropagationCollect

instance TransformationPhase ForwardPropagationCollect where
    type From ForwardPropagationCollect = ()
    type To ForwardPropagationCollect = ForwardPropagationpSemInf
    type Downwards ForwardPropagationCollect = Occurrence_place
    type Upwards ForwardPropagationCollect = (VarStatFwd, Maybe VariableData)
    downwardsBranch self d orig = occurrenceDownwards orig
    downwardsSequentialLoop self d orig = occurrenceDownwards orig
    downwardsParallelLoop self d orig = occurrenceDownwards orig
    downwardsFormalParameter self d orig = occurrenceDownwards orig
    downwardsLocalDeclaration self d orig = occurrenceDownwards orig
    downwardsAssignment self d orig = occurrenceDownwards orig
    downwardsInputActualParameter self d orig = occurrenceDownwards orig
    downwardsOutputActualParameter self d orig = occurrenceDownwards orig
    downwardsLeftValueExpression self d orig = occurrenceDownwards orig
    downwardsFunctionCall self d orig = occurrenceDownwards orig
    transformBlock self d origBlock u = Block {
        blockData = recursivelyTransformedBlockData u,
        blockSemInf = selectFromVarStatistics ( declaredVars origBlock) belowStatistics
    } where
		belowStatistics = checkFwdDeclaration (map fst $ upwardsInfoFromBlockDeclarations u) (fst $ upwardsInfoFromBlockInstructions u)
        --belowStatistics = foldl combine (fst $ upwardsInfoFromBlockInstructions u) $ map fst $ upwardsInfoFromBlockDeclarations u
    transformVariable self d origVar = origVar {
        variableSemInf = d
    }
    upwardsVariable self d origVar newVar = case d of
        Occurrence_declare  -> (Map.singleton (variableData origVar) $ Occurrences Zero Zero, Just $ variableData origVar)
        Occurrence_read -> (Map.singleton (variableData origVar) $ Occurrences Zero (One ()), Just $ variableData origVar)
        Occurrence_write ->  (Map.singleton (variableData origVar) $ Occurrences (One Nothing) Zero, Just $ variableData origVar)
        Occurrence_notopt -> (Map.singleton (variableData origVar) $ Occurrences Multiple Multiple, Just $ variableData origVar) --LIE to save variables
    upwardsSequence self d origSeq u transSeq = (checkFwdSequence $ map fst $ upwardsInfoFromSequenceProgramList u, Nothing)
    upwardsBlock self d origBlock u newBlock = (deleteFromVarStatistics (declaredVars origBlock) belowStatistics, Nothing) where
        belowStatistics = foldl combine (fst $ upwardsInfoFromBlockInstructions u) $ map fst $ upwardsInfoFromBlockDeclarations u
    upwardsParallelLoop self d origParLoop u transParLoop = (multipleVarStatistics $
        foldl combine (fst $ upwardsInfoFromParallelLoopConditionVariable u)
                    [fst $ upwardsInfoFromNumberOfIterations u, fst $ upwardsInfoFromParallelLoopCore u], Nothing)
    upwardsAssignment self d origAssign u transAssig = case assignmentLhs $ assignmentData origAssign of
        VariableLeftValue vlv -> (Map.insert var occ $ fst $ upwardsInfoFromAssignmentRhs u, Nothing)
            where
                var = variableData $ variableLeftValueContents vlv
                occ = Occurrences (One $ Just (assRs, Map.keys $ fst $ upwardsInfoFromAssignmentRhs u, False)) Zero
                assRs = case transAssig of 
                    AssignmentInstruction newAssign -> assignmentRhs $ assignmentData newAssign
                    _ -> fwdPropError $ "Internal error: ForwardPropagation/1!"
        ArrayElemReferenceLeftValue aer -> (combine (fst $ upwardsInfoFromAssignmentLhs u) (fst $ upwardsInfoFromAssignmentRhs u), Nothing)
    upwardsLocalDeclaration self d origDecl u newDecl = case  localInitValue $ localDeclarationData newDecl of
        Nothing -> defaultCase
        Just (ConstantExpression (ArrayConstant ac)) -> defaultCase
        Just initExp -> case upwardsInfoFromLocalInitValue u of
            Nothing -> defaultCase
            Just justUpFromLocalInitValue -> (Map.insert var (occ initExp $ fst justUpFromLocalInitValue) $ fst justUpFromLocalInitValue, Nothing)
        where
                    var = variableData $ localVariable $ localDeclarationData $ origDecl
                    occ initExp justUpFromLocalInitValue = Occurrences (One $ Just (initExp, Map.keys justUpFromLocalInitValue, False)) Zero
                    defaultCase = (fst $ upwardsInfoFromLocalVariable u, Nothing)
    upwardsProcedureCall self d origProcCall u transProcCall
        | List.isPrefixOf "copy" $ nameOfProcedureToCall $ procedureCallData origProcCall = case  actParams of -- TODO: eliminate string constant
            [InputActualParameter inArr, InputActualParameter arrSize, OutputActualParameter outArr] ->
                case outputActualParameterLeftValue outArr of
                    VariableLeftValue vlv -> (Map.insert (var vlv) (occ inArr) $ fst $ head ul, Nothing)
                    ArrayElemReferenceLeftValue aer -> defaultTr
            _ -> defaultTr
        | otherwise = defaultTr
        where 
            defaultTr = case ul of
                [] -> defaultValue
                otherwise -> foldl combine (head ul) (tail ul)
            ul = upwardsInfoFromActualParametersOfProcedureToCall u
            actParams = case transProcCall of
                ProcedureCallInstruction pc -> actualParametersOfProcedureToCall $ procedureCallData pc
                _ -> fwdPropError $ "Internal error: ForwardPropagation/2!"
            var vlv = variableData $ variableLeftValueContents vlv
            occ inArr = Occurrences (One $ Just (inputActualParameterExpression inArr, Map.keys $ fst $ head ul, False)) Zero
    transformSequentialLoop self d origSeqLoop u = SequentialLoopProgram $ SequentialLoop {
        sequentialLoopData = (recursivelyTransformedSequentialLoopData u) {
            conditionCalculation = (conditionCalculation $ recursivelyTransformedSequentialLoopData u) {
                blockSemInf = Map.empty 
            }
        },
        sequentialLoopSemInf = blockSemInf $ conditionCalculation $ recursivelyTransformedSequentialLoopData u
    }
    upwardsSequentialLoop self d origSeqLoop u newSeqLoop = (multipleVarStatistics $
        combine  (fst $ upwardsInfoFromSequentialLoopConditionCalculation u) $  fst $ upwardsInfoFromSequentialLoopCore u, Nothing)
    transformArrayElemReference self d origArrRef u = ArrayElemReferenceLeftValue $ ArrayElemReference {
        arrayElemReferenceData = recursivelyTransformedArrayElemReferenceData u,
        arrayElemReferenceSemInf = snd $ upwardsInfoFromArrayName u 
    }
    upwardsArrayElemReference self d origArrayRef u transArrayRefe =
        (combine (fst $ upwardsInfoFromArrayName u) (fst $ upwardsInfoFromArrayIndex u), snd $ upwardsInfoFromArrayName u)

checkFwdSequence :: [VarStatFwd]  -> VarStatFwd
checkFwdSequence [] = defaultValue
checkFwdSequence xs = List.foldl checkInSeq Map.empty xs
    where
        checkInSeq :: VarStatFwd -> VarStatFwd -> VarStatFwd
        checkInSeq preSeq curr = combine curr $ Map.mapWithKey (updatePreSeq curr) preSeq
        updatePreSeq :: VarStatFwd -> VariableData -> OccurrencesFwd -> OccurrencesFwd
        updatePreSeq curr preSeqVar preSeqOcc = case writeVar preSeqOcc of
            One (Just (preSeqExp,preSeqVars,preSeqVarsWritten))
                | preSeqVarsWritten && curr `hasRead` preSeqVar -> Occurrences (One Nothing) $ readVar preSeqOcc
                | any (hasWrite curr) preSeqVars -> case (curr `hasRead` preSeqVar)  && not ((simpleType $ variableType preSeqVar) && readVar preSeqOcc /= Multiple) of
                    True -> Occurrences (One Nothing) $ readVar preSeqOcc
                    False -> Occurrences (One (Just (preSeqExp,preSeqVars ++ (addDep curr preSeqVar),True))) $ readVar preSeqOcc
                | otherwise -> case curr `getWrite` preSeqVar of
                    Nothing -> preSeqOcc
                    Just (exp,vars,varsWritten)
                        | exp == preSeqExp -> Occurrences Zero $ readVar preSeqOcc
                        | otherwise -> preSeqOcc
            _ -> preSeqOcc
        addDep curr preSeqVar = case curr `getWrite` preSeqVar of
			Nothing -> []
			Just (exp,vars,varsWritten) -> vars

checkFwdDeclaration :: [VarStatFwd] -> VarStatFwd -> VarStatFwd
checkFwdDeclaration [] blockStat = blockStat
checkFwdDeclaration declStat blockStat = checkFwdSequence $ declStat ++ [blockStat]

-- ====================
--  ForwardPropagation
-- ====================

type VarWrite t = [(VariableData,Expression t)]

toVarWrite :: VarStatFwd -> VarWrite ForwardPropagationpSemInf
toVarWrite vs = Map.foldWithKey (getExp) [] vs where
    getExp :: VariableData -> OccurrencesFwd -> VarWrite ForwardPropagationpSemInf -> VarWrite ForwardPropagationpSemInf
    getExp name (Occurrences (One (Just (exp,_,_))) reads) vw 
        | reads /= Multiple && notConstArray exp = (name,exp):vw --used once and complex expr
        | simpleExpr exp = (name,exp):vw --used several and simple expr
        | otherwise = vw
    getExp name _ vw = vw
    notConstArray e = case e of
        ConstantExpression c -> simplConst c
        _ -> True
    simpleExpr e = case e of
        ConstantExpression c -> simplConst c
        LeftValueExpression l -> case leftValueExpressionContents l of
            VariableLeftValue v -> True
            ArrayElemReferenceLeftValue a -> simpleExpr $ arrayIndex $ arrayElemReferenceData a
        _ -> False
    simplConst (ArrayConstant ac) = False
    simplConst _ = True

instance Default (Set.Set VariableData) where
    defaultValue = Set.empty

instance Combine (Set.Set VariableData) where
    combine = Set.union

data ForwardPropagationTransform = ForwardPropagationTransform

instance TransformationPhase ForwardPropagationTransform where
    type From ForwardPropagationTransform = ForwardPropagationpSemInf
    type To ForwardPropagationTransform = ()
    type Downwards ForwardPropagationTransform = VarStatFwd
    type Upwards ForwardPropagationTransform = Set.Set VariableData
    downwardsBlock self d origBlock = combine d $ blockSemInf origBlock
    downwardsSequentialLoop self d origSeqLoop = combine d $ sequentialLoopSemInf origSeqLoop
    transformLeftValueExpression self d origLVE u = case leftValueExpressionContents origLVE of
            VariableLeftValue origVar -> case List.find (\(vn,e) -> (vn == var origVar)) varwrite of
                    Nothing -> defaultTr
                    Just repl -> fst $ walkExpression self d (snd repl)
            ArrayElemReferenceLeftValue origArr -> defaultTr
        where
            var v = variableData $ variableLeftValueContents v
            varwrite = toVarWrite d
            defaultTr = LeftValueExpression $ LeftValueInExpression {
                leftValueExpressionContents = recursivelyTransformedLeftValueExpressionContents u,
                leftValueExpressionSemInf = ()
            }
    transformVariableInLeftValue self d origVarLV u = case List.find (\(vn,e) -> (vn == var)) varwrite of
            Nothing -> defaultTr
            Just repl  -> case repl of
                    (_,LeftValueExpression lve) -> fst $ walkLeftValue self d $ leftValueExpressionContents lve
                    _ -> defaultTr
        where 
            var = variableData $ variableLeftValueContents origVarLV
            varwrite = toVarWrite d
            defaultTr = VariableLeftValue $ VariableInLeftValue {
                variableLeftValueContents = recursivelyTransformedVariableLeftValueContents u,
                variableLeftValueSemInf = ()
            }
    transformArrayElemReference self d origArrayRef u = case List.find (\(vn,e) -> (vn == var)) varwrite of
            Nothing -> defaultTr
            Just repl  -> case repl of
                    (_,LeftValueExpression lve) -> case leftValueExpressionContents lve of
                        VariableLeftValue vlv -> defaultTr
                        ArrayElemReferenceLeftValue aer -> ArrayElemReferenceLeftValue $ ArrayElemReference {
                            arrayElemReferenceData = ArrayElemReferenceData { 
                                arrayName = fst $ walkLeftValue self (newD d var aer origArrayRef) $ arrayName $ arrayElemReferenceData origArrayRef
                                , arrayIndex = fst $ walkExpression self d $ arrayIndex $ arrayElemReferenceData aer
                            },
                            arrayElemReferenceSemInf = ()
                        }
                    _ -> defaultTr
        where
            newD :: VarStatFwd -> VariableData -> ArrayElemReference ForwardPropagationpSemInf -> ArrayElemReference ForwardPropagationpSemInf -> VarStatFwd
            newD d var rep orig = Map.adjust (newDD var rep orig) var d
            newDD var rep orig x = x {
                writeVar = One $ Just ( LeftValueExpression $ LeftValueInExpression {
                    leftValueExpressionContents = ArrayElemReferenceLeftValue $ ArrayElemReference {
                        arrayElemReferenceData = ArrayElemReferenceData { 
                            arrayName = arrayName $ arrayElemReferenceData rep
                            , arrayIndex = arrayIndex $ arrayElemReferenceData orig
                        },
                        arrayElemReferenceSemInf = Just  var
                    },  
                    leftValueExpressionSemInf = () 
                },[],False)
            }
            var = getJust $ arrayElemReferenceSemInf origArrayRef
            getJust (Just a) = a
            varwrite = toVarWrite d
            defaultTr = ArrayElemReferenceLeftValue $ ArrayElemReference {
                arrayElemReferenceData = recursivelyTransformedArrayElemReferenceData u,
                arrayElemReferenceSemInf = convert $ arrayElemReferenceSemInf origArrayRef
            }
    upwardsVariable self d origVar newVar = case variableSemInf origVar of
        Occurrence_declare  -> Set.empty
        Occurrence_read -> Set.empty
        Occurrence_write -> Set.singleton (variableData origVar)
        Occurrence_notopt -> Set.empty
    upwardsBlock self d origBlock u transformedBlock = foldl (\s e -> Set.delete e s) (upwardsInfoFromBlockInstructions u) (declaredVars origBlock) --Not need just optimalize compliler (not try delete locals outside block)
    transformBlock self d origBlock u = delUnusedDecl (map fst $ toVarWrite $ combine d $ blockSemInf origBlock) origBlock (recursivelyTransformedBlockData u)
    transformPrimitive self d originalPrimitive u = case canDelete of
            True -> EmptyProgram $ Empty ()
            False ->
                PrimitiveProgram $ Primitive {
                    primitiveInstruction = recursivelyTransformedPrimitiveInstruction u,
                    primitiveSemInf = ()
                }
        where
            canDelete = Set.isSubsetOf (upwardsInfoFromPrimitiveInstruction u) (Set.fromList $ map fst $ toVarWrite d)