Unroll.hs

{-
 - Copyright (c) 2009, 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.
 -}

module Feldspar.Compiler.Optimization.Unroll where

import Feldspar.Compiler.Imperative.Representation hiding (None)
import Feldspar.Compiler.Options
import Feldspar.Compiler.Optimization.Replace
import Prelude

-- | Unroll opreation for imperative functions.
doUnroll :: Options -> [ImpFunction] -> [ImpFunction]
doUnroll opt ps = map (doUnrollOne opt) ps

-- Unroll opreation for an Imperative function.
doUnrollOne :: Options -> ImpFunction -> ImpFunction
doUnrollOne opt p = case unroll opt of 
    NoUnroll -> p
    (Unroll i) -> unrollStruc i p

-- If the second parameter is a For loop which contains Empty, Primitive or Seq programtypes and the modulo of the first parameret and the maximum iteration of loop, will terurn true.
unrollPossible :: Int -> Program -> Bool
unrollPossible i (ParLoop counter num 1 (CompPrg _ prg) inf) = moduloOk && unrollPossible' prg
    where
        moduloOk = case num of 
            Expr (ConstExpr (IntConst x)) _     -> x `mod` i == 0
            otherwise                           -> True
        unrollPossible' Empty = True
        unrollPossible' (Primitive i s) = True
        unrollPossible' (Seq ps si) = and $ map unrollPossible' ps
        unrollPossible' _ = False
unrollPossible _ _ = False

-- Collects variable names from a declaration bloc.
collectVars :: [Declaration] -> [String]
collectVars ds = map collectVar ds where
    collectVar (Decl (Var s _ _) declType initVal inf) = s

-- Concatenates the first and second parameters and returnes as a string.
alterVarName :: String -> Int -> String
alterVarName old idx = old ++ "_" ++ show idx

-- Creates a new additional expression from the loop counter and a positive constant.
alterVar :: String -> Int -> UntypedExpression
alterVar name idx = FunCall InfixOp "+" [var,const]
    where
        var = Expr (LeftExpr $ LVar $ Var name Normal int) int
        const = Expr (ConstExpr $ IntConst idx) int
        int = (Numeric ImpSigned S32)
 
-- Replicates the declarated variables with new names.
unrollDecl :: [Declaration] -> String -> Int -> [Declaration]
unrollDecl decllist loopvar i 
    = unrollDecl' decllist todolists
    where 
        todolists = zip (replicate i ((collectVars decllist),loopvar)) [0,1..(i-1)]
        unrollDecl' :: [Declaration] -> [(([String],String),Int)] -> [Declaration]
        unrollDecl' decllist todolists = foldl (++) [] (map (unrollOneDecl decllist) todolists) 
        unrollOneDecl:: [Declaration] -> (([String],String),Int) -> [Declaration]
        unrollOneDecl decllist ((local_vars,loopvar),idx) 
            | idx < 1   = foldl (\decllist var -> (replaceVar decllist (var,alterVarName var idx))) decllist local_vars
            | otherwise = foldl (\decllist var -> (replaceVar decllist (var,alterVarName var idx))) 
                            (replaceUExpr decllist (loopvar,(alterVar loopvar idx))) local_vars 

-- Replicates the for loop body using the new variables.
unrollPrg :: Program -> String -> [String] -> Int -> [Program]
unrollPrg prg loopvar locals num =
    map alter $ zip (replicate num prg) [0..] where
        alter (p,idx) 
            | idx < 1 =  foldl (\p' tr -> tr p') p (map (alterLocal idx) locals)
            | otherwise = foldl (\p' tr -> tr p') (alterLoopVar (p,idx)) (map (alterLocal idx) locals)
        alterLoopVar (p,idx) = replaceUExpr p (loopvar, alterVar loopvar idx)
        alterLocal idx loc p = replaceVar p (loc, alterVarName loc idx)

-- Unrolls the declaration and the body of the loop, if the unroll operation is possible. If not possible, then tries to find sub-loops in the current loop.
unrollRepeatSimple :: Program -> Int -> Program
unrollRepeatSimple p i = urs p i (unrollPossible i p) where
    urs (ParLoop (Var v k t) max step cprg inf) i True
        = ParLoop (Var v k t) max (step*i) (CompPrg (unrollDecl (locals cprg) v i) (Seq (unrollPrg (body cprg) v (collectVars (locals cprg)) i) inf)) inf
    urs (ParLoop (Var v k t) max step cprg inf) i False
        = ParLoop (Var v k t) max step cprg{ body = unrollStruc i (body cprg)} inf
    urs p i False = p

-- Finds for loops in data hierarchy and make the unroll opertaion.



class Unroll t where
    unrollStruc :: Int -> t -> t

instance Unroll ImpFunction where
    unrollStruc i f = f{ prg = unrollStruc i $ prg f }

instance Unroll CompleteProgram where
    unrollStruc i c = c{ body = unrollStruc i $ body c }

instance Unroll Program where
    unrollStruc i (Seq ps inf) = Seq (map (unrollStruc i) ps) inf
    unrollStruc i (IfThenElse v cpt cpe inf) = IfThenElse v (unrollStruc i cpt) (unrollStruc i cpe) inf
    unrollStruc i for@(ParLoop _ _ _ _ _) = unrollRepeatSimple for i
    unrollStruc i (SeqLoop v calc body inf) = SeqLoop v (unrollStruc i calc) (unrollStruc i body) inf
    unrollStruc _ x = x