-----------------------------------------------------------------------------
-- |
-- 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.MShow
import WinDll.Structs.MShow.HaskellSrcExts
import WinDll.Structs.Folds.HaskellSrcExts
import WinDll.Utils.Types
import qualified Debug.Trace as D

-- | 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 (show t') $ D.trace ("IN:  " ++ mshowM 2 t') $ D.trace ("OUT: " ++ mshowM 2 ty') $ D.trace (show arr' ++ " - " ++ show lst ++ " (" ++ show arr' ++ ")") 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 
                                    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                = g x
        
        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
                                                
        -- | 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
                        
-- | Updates the type of the higher-order function to be in IO
hofIO :: Exts.Type -> Exts.Type
hofIO = everywhere (mkT hof)
  where hof :: Exts.Type -> Exts.Type
        hof (Exts.TyParen ty) = Exts.TyParen (mkIO ty)
        hof ty                = 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'}