-----------------------------------------------------------------------------
-- |
-- Module      :  Windll
-- Copyright   :  (c) Tamar Christina 2009 - 2010
-- License     :  BSD3
-- 
-- Maintainer  :  tamar@zhox.com
-- Stability   :  experimental
-- Portability :  portable
--
-- Various functions to detect and record list types within a 
-- given type.
--
-----------------------------------------------------------------------------

module WinDll.Utils.ListTypes where

import qualified Language.Haskell.Exts as Exts

import Data.Generics
import Data.List

import WinDll.Structs.Structures hiding (Module)
import qualified WinDll.Structs.Structures as WinDll
import WinDll.Structs.MShow.HaskellSrcExts
import WinDll.Structs.Folds.HaskellSrcExts
import WinDll.Utils.Types

-- | Upgrade a type by performing actions such as identifying list 
--   and changing the type to also pass along list counters
upgradeType :: Ann -> Exts.Type -> (Ann, Exts.Type)
upgradeType ann x = 
    let t'   = updateType True ann' t
        t    = simplify x
        lst' = findListIndices ty'
        ann' = ann{annArrayIndices = lst'}
        ty'  = analyzeType True t'
    in (ann', ty')
    
-- | Check to see if the last argument of the type is a list
--   in which case we should change the Int before it to Ptr CInt
analyzeType :: Bool -> Exts.Type -> Exts.Type
analyzeType esc t =
   let t'   = everywhere (mkT embedded) t
       lst  = findListIndices t'
       arr' = tlength t' - 1
       part = arr' `elem` lst 
       ty'  = if part 
                 then changeType esc arr' (\val -> case val of
                                                      Exts.TyCon (Exts.UnQual (Exts.Ident "Int"))  -> (Exts.TyApp 
                                                            (Exts.TyCon $ Exts.UnQual $ Exts.Ident "Ptr")
                                                            (Exts.TyCon $ Exts.UnQual $ Exts.Ident $ if esc then "CInt" else "Int"))
                                                      _                                            -> val) t'
                 else t'
   in ty' -- D.trace ("IN:  " ++ mshowM 2 t') $ D.trace ("OUT: " ++ mshowM 2 ty') $ D.trace (show arr' ++ " - " ++ show lst) ty'
  where embedded :: Exts.Type -> Exts.Type
        embedded (Exts.TyParen a) = (Exts.TyParen (analyzeType esc a))
        embedded x                = x
  
-- | Update the n-th element of the type with whatever we want
changeType :: Bool -> Int -> (Exts.Type -> Exts.Type) -> Exts.Type -> Exts.Type
changeType _esc n f ty = fst $ (foldTypeIO alg ty) 0
  where alg :: TypeAlgebraIO (Int -> (Exts.Type, Int))
        alg = (\a b c i -> let (c', i') = c i
                           in (Exts.TyForall a b c', i')
              ,\a b   i -> let (a', i' ) = a i
                               (b', i'') = b i'
                           in (Exts.TyFun a' b', i'')
              ,\a b   i -> let (b', i') = app b i
                           in (Exts.TyTuple a b', i)
              ,\a     i -> let (a', i') = a i
                           in (Exts.TyList a', i')
              ,\o a b i -> let (a', i' ) = a i
                               (b', i'') = b i'
                               ix        = if o then  i'' else i'
                           in (Exts.TyApp a' b', ix) -- i'')
              ,\a     i -> let i' = i + 1
                               a' = Exts.TyVar a
                           in if i' == n 
                                 then (f a', i')
                                 else (a'  , i')
              ,\a     i -> let i' = i + 1
                               a' = Exts.TyCon a
                           in if i' == n 
                                 then (f a', i')
                                 else (a'  , i')
              ,\a     i -> let (a', i') = a i
                           in (Exts.TyParen a', i')
              ,\a b c i -> let (a', i' ) = a i
                               (c', i'') = c i'
                           in (Exts.TyInfix a' b c', i'')
              ,\a b   i -> let (a', i') = a i
                           in (Exts.TyKind a' b, i')
              )
        app :: [Int -> (Exts.Type, Int)] -> Int -> ([Exts.Type], Int)
        app []     i = ([], i)
        app (x:xs) i = let (x' , i' ) = x i
                           (xs', i'') = app xs i'
                       in (x':xs', i'')
    
-- | Update a type according to the annotations present
updateType :: Bool -> Ann -> Exts.Type -> Exts.Type
updateType esc ann = everywhere (mkT pushType)
  where -- | Types to update
        pushType :: Exts.Type -> Exts.Type
        pushType (Exts.TyFun a b) = let f x = case isIOList x of
                                                True  -> Exts.TyFun
                                                          (Exts.TyCon $ Exts.UnQual $ Exts.Ident "Int")
                                                False -> id
                                        g x = case isIOList x of
                                                True  -> Exts.TyFun 
                                                            (Exts.TyApp 
                                                                (Exts.TyCon $ Exts.UnQual $ Exts.Ident "Ptr")
                                                                (Exts.TyCon $ Exts.UnQual $ Exts.Ident $ if esc then "CInt" else "Int"))
                                                             x
                                                False -> x
                                    in f a $ Exts.TyFun a (g b)
        -- pushType (Exts.TyApp a b) = let f x = case isList x of
                                                -- True  -> Exts.TyFun
                                                          -- (Exts.TyCon $ Exts.UnQual $ Exts.Ident "Int") x
                                                -- False -> x
                                     -- in if isIO a 
                                           -- then simplify (move a $ f b)
                                           -- else Exts.TyApp a b
        pushType x                = x
        
        -- | Move an IO declaration inwards.
        move :: Exts.Type -> Exts.Type -> Exts.Type
        move io (Exts.TyFun a b) = Exts.TyFun a (Exts.TyApp io b)
        move io rest             = Exts.TyApp io rest


-- | Identifies locations within a Type where lists are found
--   The indices provided are the locations of the size variables
--   of arrays. The counters start at 0 and not 1 anymore.
--   So keep this in mind :)
findListIndices :: Exts.Type -> [Int]
findListIndices ty = coords (embed ty) 
  where embed :: Exts.Type -> [Bool]
        embed (Exts.TyFun a b) = let res = isList a
                                 in if isFun b
                                       then res:embed b
                                       else res:[isIOList b]
        embed (Exts.TyParen a) = embed a
        embed _ = []
        
        coords :: [Bool] -> [Int]
        coords b = [i | (x,i) <- zip b [(-1)..], x]
        
-- | A variant of isList that looks inside IO
isIOList :: Exts.Type -> Bool
isIOList (Exts.TyApp a b) = if isIO a 
                               then isList b
                               else False
isIOList x                = isList x
                       
-- | Update the n-th element of the type with whatever we want,
--   only looking at the amount of (->) constructors.
processTypeNode :: Int -> (Exts.Type -> Exts.Type) -> Exts.Type -> Exts.Type
processTypeNode n f ty = fst $ (foldTypeIO alg ty) 1
  where alg :: TypeAlgebraIO (Int -> (Exts.Type, Int))
        alg = (\a b c i -> let (c', i') = c i
                           in (Exts.TyForall a b c', i')
              ,\a b   i -> let (a', i' ) = a i
                               (b', i'') = b (i + 1)
                               value     = Exts.TyFun (if i   == n then f a' else a')
                                                      (if i+1 == n then f b' else b')
                           in (value , i'')
              ,\a b   i -> let (b', i') = app b i
                           in (Exts.TyTuple a b', i)
              ,\a     i -> let (a', i') = a i
                           in (Exts.TyList a', i')
              ,\o a b i -> let (a', i' ) = a 0
                               (b', i'') = b 0
                           in (Exts.TyApp a' b', i)
              ,\a     i -> (Exts.TyVar a, i)
              ,\a     i -> (Exts.TyCon a, i)
              ,\a     i -> let (a', i') = a i
                           in (Exts.TyParen a', i')
              ,\a b c i -> let (a', i' ) = a i
                               (c', i'') = c i'
                           in (Exts.TyInfix a' b c', i'')
              ,\a b   i -> let (a', i') = a i
                           in (Exts.TyKind a' b, i')
              )
        app :: [Int -> (Exts.Type, Int)] -> Int -> ([Exts.Type], Int)
        app []     i = ([], i)
        app (x:xs) i = let (x' , i' ) = x i
                           (xs', i'') = app xs i'
                       in (x':xs', i'')

-- | Updates a type to that which uses IO
mkIO :: Exts.Type -> Exts.Type
mkIO ty = let arr = tlength ty
              mk  = simplify . Exts.TyApp (Exts.TyCon $ Exts.UnQual $ Exts.Ident "IO")
          in if isIO ty 
                then ty
                else if arr == 1 -- if there are no arguments, just directly apply mk
                        then mk ty
                        else processTypeNode arr mk ty

-- | Checks to see if the function being returned is in IO
isIO :: Exts.Type -> Bool
isIO ty = let tys = collectLessTypes ty
              ret = last tys
          in "IO" `isPrefixOf` ret
                              
-- | See if the type is just a list type
isOnlyList :: Exts.Type -> Bool
isOnlyList (Exts.TyParen a) = isOnlyList a
isOnlyList (Exts.TyList  _) = True
isOnlyList _                = False

updateModule :: WinDll.Module -> WinDll.Module
updateModule = everywhere (mkT mkFunction `extT` mkExport `extT` mkDataType)
  where mkFunction e@(WinDll.Function{}) = let (newAnn, ty) = upgradeType (fnAnn e) (fnType e)
                                           in e{fnType = ty, fnAnn = newAnn }
        mkExport   e@(WinDll.Export{}  ) = let (_, ty) = upgradeType noAnn (exType e)
                                           in e{exType = ty}
        mkDataType e@(WinDll.DataType{}) = let dt = dtCons e
                                           in e{dtCons = map mkConstr dt}
        mkDataType e@(WinDll.NewType{} ) = let dt = dtCon e
                                           in e{dtCon = mkConstr dt}
        mkDataType e                     = e
        mkConstr   e@(WinDll.Constr{}  ) = let dt = dtNamed e
                                           in e{dtNamed = map mkAnnNamedTypes dt}
        mkAnnNamedTypes e@(AnnType{}   ) = let (newAnn, ty) = upgradeType (antAnn e) (antType e)
                                               ann'         = newAnn{annArrayIsList = True, annArrayIndices = []}
                                           in case isOnlyList (antType e) of
                                                False -> e{antType = ty, antAnn = newAnn}
                                                True  -> e{antAnn = ann'}