module Feldspar.Compiler.Backend.C.Plugin.AllocationEliminator where

import Data.Map (Map)
import qualified Data.Map as Map
import Feldspar.Transformation

data AllocationEliminator = AllocationEliminator

instance Transformation AllocationEliminator
  where
    type From AllocationEliminator = ()
    type To AllocationEliminator = ()
    type Down AllocationEliminator = ()
    type Up AllocationEliminator = ()
    type State AllocationEliminator = (Integer, Map String Integer)

instance Transformable AllocationEliminator Definition
  where
    transform t s d proc@(Procedure _ _ _ _ _ _) = Result proc'{ inParams = mem : inParams proc' } s' u'
      where
        Result proc' s' u' = defaultTransform t s d proc
        mem = Variable
            { varName = "mem"
            , varType = ArrayType UndefinedLen $ ArrayType UndefinedLen VoidType
            , varRole = Value
            , varLabel = ()
            }

    transform t s d x = defaultTransform t s d x

instance Transformable AllocationEliminator Expression
  where
    transform t s@(idx,m) d e@(VarExpr v lab) = case Map.lookup (varName v) m of
        Nothing -> defaultTransform t s d e
        Just i  -> Result ArrayElem
            { array         = VarExpr
                { var   = Variable
                    { varName = "mem"
                    , varType = ArrayType UndefinedLen $ varType v
                    , varRole = Value
                    , varLabel = ()
                    }
                , exprLabel = ()
                }
            , arrayIndex    = ConstExpr
                { constExpr = IntConst
                    { intValue = i
                    , intConstLabel = ()
                    , constLabel = ()
                    }
                , exprLabel = ()
                }
            , arrayLabel = ()
            , exprLabel = ()
            } s ()
    transform t s d e = defaultTransform t s d e

instance Transformable1 AllocationEliminator [] Declaration
  where
    transform1 t s d [] = Result1 [] s ()
    transform1 t s@(idx,m) d (x:xs) = case varType $ declVar x of
        ArrayType _ _   -> transform1 t (idx+1, Map.insert (varName $ declVar x) idx m) d xs
        _               -> Result1 (x:xs') s' u'
      where
        Result1 xs' s' u' = transform1 t s d xs

instance Plugin AllocationEliminator
  where
    type ExternalInfo AllocationEliminator = ()
    executePlugin self@AllocationEliminator externalInfo procedure = 
        result $ transform self (0,Map.empty) () procedure