{-
 - 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 FlexibleInstances, TypeSynonymInstances #-}

module Feldspar.Compiler.Transformation.GraphUtils
 ( tupleWalk
 , tupleZip
 , tupleZipList
 , replaceVars
 ) where

import Feldspar.Core.Graph
import Feldspar.Core.Types
import Data.List


-- replaceVars [(var,fun)] ndhier ---- replace the variable (or the variables of the same node 
----- if the list part is empty) according to the "fun"
class RepVars a where
   replaceVars:: [(Variable, Variable -> Variable)] -> a ->  a
    
instance RepVars (Node, [Hierarchy]) where
   replaceVars chLs (node, hs) =  (replaceVars chLs node, map (replaceVars chLs) hs)

instance RepVars Hierarchy where
   replaceVars chLs (Hierarchy ndHrs) = Hierarchy (map (replaceVars chLs) ndHrs)

instance RepVars Node where
   replaceVars chLs (node@(Node {input = nInp, function = nFunc})) 
           = node{input= replaceVars chLs nInp, function = replaceVars chLs nFunc}
   
instance RepVars (Tuple Source) where
   replaceVars chLs (One (Constant x)) = One (Constant x)
   replaceVars chLs (One (Variable x)) = One (Variable (replaceVars chLs x))
   replaceVars chLs (Tup tls) = Tup (map (replaceVars chLs) tls)
   
instance RepVars Variable where
   replaceVars chLs (nId, ls) 
       = case find (\((v,_),_) -> v == nId) chLs of
           Nothing -> (nId, ls)
           Just ((v,vls),tr) -> case vls of
                                     [] -> tr (nId, ls)
                                     _  -> if (vls == ls) then (tr (nId,ls)) else (nId,ls)    
            
instance RepVars Function where
   replaceVars chLs (NoInline str ifc) = (NoInline str (replaceVars chLs ifc))
   replaceVars chLs (Parallel ifc) = (Parallel (replaceVars chLs ifc))
   replaceVars chLs (IfThenElse ifc1 ifc2) = (IfThenElse (replaceVars chLs ifc1) (replaceVars chLs ifc2))
   replaceVars chLs (While ifc1 ifc2) = (While (replaceVars chLs ifc1) (replaceVars chLs ifc2))
   replaceVars chLs fun = fun

instance RepVars Interface where
   replaceVars chLs ifc@ (Interface {interfaceOutput = ifOut}) 
           = ifc{interfaceOutput = replaceVars chLs ifOut}

-- The 'tupleWalk' function walks through a tuple, applies the given
-- function to every leaf (while provides information about the place of
-- the leaf) and puts the results in a list.
tupleWalk :: ([Int] -> a -> b) -> Tuple a -> [b]
tupleWalk = tupleWalk' [] where
    tupleWalk' :: [Int] -> ([Int] -> a -> b) -> Tuple a -> [b]
    tupleWalk' p f (One x) = [f p x]
    tupleWalk' p f (Tup xs) = concatMap ff $ zip xs [0..] where
        ff (x,idx) = tupleWalk' (p ++ [idx]) f x

-- Zips to tuples of the same structure.
tupleZip :: (Tuple a, Tuple b) -> Tuple (a,b)
tupleZip (One x, One y) = One (x,y)
tupleZip (Tup xs, Tup ys) = Tup (map tupleZip $ zip xs ys)
tupleZip _ = error "Error: Tuples with different structure are zipped."

-- Zips the "leafs" to list of tuples.
tupleZipList :: (Tuple a, Tuple b) -> [(a,b)]
tupleZipList (One x, One y) = [(x,y)]
tupleZipList (Tup xs, Tup ys) = concatMap tupleZipList $ zip xs ys
tupleZipList _ = error "Error: Tuples with different structure are zipped."