Precompilation.hs

{-# LANGUAGE EmptyDataDecls, TypeFamilies #-}

module Feldspar.Compiler.Plugins.Precompilation where

import Feldspar.Compiler.PluginArchitecture
import qualified Feldspar.Core.Expr as Expr
import Feldspar.Core.Types

import qualified Feldspar.Compiler.Precompiler.Precompiler as Precompiler
import Feldspar.Compiler.Error

import System.IO.Unsafe

-- ===========================================================================
--  == Precompilation plugin
-- ===========================================================================

data CompilationMode = Interactive | Standalone
    deriving (Show, Eq)

data SignatureInformation = SignatureInformation {
    originalFeldsparFunctionName      :: String,
    generatedImperativeParameterNames :: [String],
    originalFeldsparParameterNames    :: Maybe [Maybe String]
} deriving (Show, Eq)

instance Default SignatureInformation where defaultValue = precompilationError InternalError "Default value should not be used"

precompilationError = handleError "PluginArch/Precompilation"

data PrecompilationSemanticInfo

instance SemanticInfo PrecompilationSemanticInfo where
    type ProcedureInfo PrecompilationSemanticInfo = SignatureInformation
    type BlockInfo PrecompilationSemanticInfo = ()
    type ProgramInfo PrecompilationSemanticInfo = ()
    type EmptyInfo PrecompilationSemanticInfo = ()
    type PrimitiveInfo PrecompilationSemanticInfo = ()
    type SequenceInfo PrecompilationSemanticInfo = ()
    type BranchInfo PrecompilationSemanticInfo = ()
    type SequentialLoopInfo PrecompilationSemanticInfo = ()
    type ParallelLoopInfo PrecompilationSemanticInfo = ()
    type FormalParameterInfo PrecompilationSemanticInfo = ()
    type LocalDeclarationInfo PrecompilationSemanticInfo = ()
    type LeftValueExpressionInfo PrecompilationSemanticInfo = ()
    type VariableInLeftValueInfo PrecompilationSemanticInfo = ()
    type ArrayElemReferenceInfo PrecompilationSemanticInfo = ()
    type InputActualParameterInfo PrecompilationSemanticInfo = ()
    type OutputActualParameterInfo PrecompilationSemanticInfo = ()
    type AssignmentInfo PrecompilationSemanticInfo = ()
    type ProcedureCallInfo PrecompilationSemanticInfo = ()
    type FunctionCallInfo PrecompilationSemanticInfo = ()
    type IntConstantInfo PrecompilationSemanticInfo = ()
    type FloatConstantInfo PrecompilationSemanticInfo = ()
    type BoolConstantInfo PrecompilationSemanticInfo = ()
    type ArrayConstantInfo PrecompilationSemanticInfo = ()
    type VariableInfo PrecompilationSemanticInfo = SignatureInformation

data Precompilation = Precompilation

instance TransformationPhase Precompilation where
    type From Precompilation = ()
    type To Precompilation = ()
    type Downwards Precompilation = SignatureInformation
    type Upwards Precompilation = ()
    downwardsProcedure Precompilation fromAbove procedure = fromAbove {
        generatedImperativeParameterNames =
            map (variableName . variableData . formalParameterVariable) (inParameters procedure)
    }
    transformProcedure Precompilation fromAbove originalProcedure fromBelow =
        Procedure { -- NOTE: fromAbove won't have the generated imperative parameter names right here
            procedureName = originalFeldsparFunctionName fromAbove,
            inParameters  = recursivelyTransformedInParameters fromBelow,
            outParameters = recursivelyTransformedOutParameters fromBelow,
            procedureBody = recursivelyTransformedProcedureBody fromBelow,
            procedureSemInf = ()
        }
    transformVariable = myTransformVariable

getVariableName :: SignatureInformation -> String -> String
getVariableName signatureInformation origname = case (originalFeldsparParameterNames signatureInformation) of
    Just originalParameterNameList ->
        if length (generatedImperativeParameterNames signatureInformation) == length originalParameterNameList then
            case searchResults of
                [] -> origname
                otherwise -> case snd $ head $ searchResults of
                                Just newname -> newname
                                Nothing -> origname
        else
            precompilationError InternalError $ "parameter name list length mismatch:" ++
                    show (generatedImperativeParameterNames signatureInformation) ++ " " ++ show originalParameterNameList
        where
            searchResults = (filter (((==) origname).fst) (zip (generatedImperativeParameterNames signatureInformation) originalParameterNameList))
    Nothing -> origname

myTransformVariable :: Precompilation -> SignatureInformation -> Variable () -> Variable ()
myTransformVariable Precompilation fromAbove v = v {
    variableData = (variableData v) {
        variableName = getVariableName fromAbove (variableName $ variableData v)
    },
    variableSemInf = ()
}

data PrecompilationExternalInfo = PrecompilationExternalInfo {
    originalFeldsparFunctionSignature :: Precompiler.OriginalFeldsparFunctionSignature, 
    graphInputInterfaceType :: Tuple StorableType,
    numberOfFunctionArguments :: Int,
    compilationMode :: CompilationMode
}

countTuple :: Tuple a -> Int
countTuple (One x) = 1
countTuple (Tup list) = sum (map countTuple list)

addPostfixNumbersToMaybeList :: [Maybe String] -> [Maybe String]
addPostfixNumbersToMaybeList list
    | length list > 1 = map addPostfixNumberToMaybeString (zip list [1..]) -- postfix numbers only needed for lists with length > 1
    | otherwise = list

addPostfixNumberToMaybeString :: (Maybe String, Int) -> Maybe String
addPostfixNumberToMaybeString (ms, num) = case ms of
    Just s -> Just $ s ++ (show num)
    Nothing -> Nothing
    
inflate :: Int -> [Maybe String] -> [Maybe String]
inflate target list | length list < target = inflate target (list++[Nothing])
                    | length list == target = list
                    | otherwise = precompilationError InternalError "Unexpected situation in 'inflate'"
    
-- Applies some tweaks the original parameter name list based on the graph's input interface type signature
parameterNameListConsolidator :: PrecompilationExternalInfo -> [Maybe String]
parameterNameListConsolidator externalInfo = case graphInputInterfaceType externalInfo of
    One x -> Precompiler.originalFeldsparParameterNames $ originalFeldsparFunctionSignature externalInfo
    tuple@(Tup list) -> case numberOfFunctionArguments externalInfo of
        0 -> precompilationError InternalError "parameter name list consolidator function shouldn't be called when numArgs==0"
        1 -> addPostfixNumbersToMaybeList $ replicate (countTuple tuple)
                 (head $ Precompiler.originalFeldsparParameterNames $ originalFeldsparFunctionSignature externalInfo)
        otherwise -> concat $ map (\(cnt,name)->addPostfixNumbersToMaybeList (replicate cnt name)) 
           (zip (map countTuple list) (Precompiler.originalFeldsparParameterNames $ originalFeldsparFunctionSignature externalInfo))

instance Plugin Precompilation where
    type ExternalInfo Precompilation = PrecompilationExternalInfo
    executePlugin Precompilation externalInfo procedure = fst
        $ executeTransformationPhase Precompilation (SignatureInformation {
            originalFeldsparFunctionName = Precompiler.originalFeldsparFunctionName $ originalFeldsparFunctionSignature externalInfo,
            generatedImperativeParameterNames = precompilationError InternalError "GIPN should have been overwritten", 
            originalFeldsparParameterNames = if numberOfFunctionArguments externalInfo == 0
                then
                    Nothing -- if there are no arguments, disable parameter name handling (needed because of the dummy var0)
                else
                    (case compilationMode externalInfo of
                        Standalone ->
                            if -- ultimate check, should be enough...
                                numberOfFunctionArguments externalInfo ==
                                length (Precompiler.originalFeldsparParameterNames $ originalFeldsparFunctionSignature externalInfo)
                            then
                                Just $ parameterNameListConsolidator externalInfo
                            else
                                (unsafePerformIO $ do
                                    putStrLn $ "[WARNING @ PluginArch/Precompilation]: argument count mismatch in function " ++ 
                                          (Precompiler.originalFeldsparFunctionName $ originalFeldsparFunctionSignature externalInfo) ++
                                          ", inflating incomplete parameter name list..."
                                    putStrLn $ "numArgs: " ++ show (numberOfFunctionArguments externalInfo) ++ ", parameter list: " ++ 
                                        show (Precompiler.originalFeldsparParameterNames $ originalFeldsparFunctionSignature externalInfo) 
                                    return $ Just $ parameterNameListConsolidator (externalInfo {
                                        originalFeldsparFunctionSignature = (originalFeldsparFunctionSignature externalInfo) {
                                            Precompiler.originalFeldsparParameterNames = inflate (numberOfFunctionArguments externalInfo) $
                                                Precompiler.originalFeldsparParameterNames $ originalFeldsparFunctionSignature externalInfo
                                        }
                                    })
                                )
                        Interactive -> Nothing -- no parameter name handling in interactive mode
                    )
         }) procedure