content: {-# LANGUAGE GADTs               #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  -- |
  -- Module      : Data.Array.Accelerate.Pretty.Traverse
  -- Copyright   : [2010..2011] The Accelerate Team
  -- License     : BSD3
  --
  -- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
  -- Stability   : experimental
  -- Portability : non-portable (GHC extensions)
  --
  module Data.Array.Accelerate.Pretty.Traverse
    where
  
  -- friends
  import Data.Array.Accelerate.AST
  import Data.Array.Accelerate.Array.Representation ( SliceIndex (..) )
  import Data.Array.Accelerate.Array.Sugar          hiding ((!))
  import Data.Array.Accelerate.Type
  
  
  cat :: Show s => String -> s -> String
  cat t s = t ++ show s
  
  data Labels = Labels { accFormat :: String
                       , expFormat :: String
                       , funFormat :: String
                       , tupleFormat :: String
                       , arrayFormat :: String
                       , boundaryFormat :: String
                       , primFunFormat :: String
                       }
  
  travAcc :: forall m b aenv a. Monad m => Labels -> (String -> String -> [m b] -> m b)
         -> (String -> String -> m b) -> OpenAcc aenv a -> m b
  travAcc f c l (OpenAcc openAcc) = travAcc' openAcc
    where
      combine = c (accFormat f)
      leaf    = l (accFormat f)
  
      travAcc' :: PreOpenAcc OpenAcc aenv a -> m b
      travAcc' (Alet acc1 acc2)                      = combine "Alet" [travAcc f c l acc1, travAcc f c l acc2]
      travAcc' (Avar idx)                            = leaf   ("AVar " `cat` idxToInt idx)
      travAcc' (Apply afun acc)                      = combine "Apply" [travAfun f c l afun, travAcc f c l acc]
      travAcc' (Aforeign ff afun acc)                = combine ("Aforeign " ++ strForeign ff) [travAfun f c l afun, travAcc f c l acc]
      travAcc' (Acond e acc1 acc2)                   = combine "Acond" [travExp f c l e, travAcc f c l acc1, travAcc f c l acc2]
      travAcc' (Awhile cond body acc)                = combine "Awhile" [travAfun f c l cond, travAfun f c l body, travAcc f c l acc]
      travAcc' (Atuple tup)                          = combine "Atuple" [ travAtuple f c l tup ]
      travAcc' (Aprj idx a)                          = combine ("Aprj " `cat` tupleIdxToInt idx) [ travAcc f c l a ]
      travAcc' (Use arr)                             = combine "Use" [ travArrays f c l (arrays (undefined::a)) arr ]
      travAcc' (Unit e)                              = combine "Unit" [ travExp f c l e ]
      travAcc' (Generate sh fun)                     = combine "Generate" [ travExp f c l  sh, travFun f c l fun]
      travAcc' (Transform sh pf vf acc)              = combine "Transform" [ travExp f c l sh, travFun f c l pf, travFun f c l vf, travAcc f c l acc ]
      travAcc' (Reshape sh acc)                      = combine "Reshape" [ travExp f c l  sh, travAcc f c l acc ]
      travAcc' (Replicate _ ix acc)                  = combine "Replicate" [ travExp f c l  ix, travAcc f c l acc ]
      travAcc' (Slice _ acc ix)                      = combine "Slice" [ travAcc f c l acc, travExp f c l  ix ]
      travAcc' (Map fun acc)                         = combine "Map" [ travFun f c l fun, travAcc f c l acc ]
      travAcc' (ZipWith fun acc1 acc2)               = combine "ZipWith" [ travFun f c l fun, travAcc f c l acc1, travAcc f c l acc2 ]
      travAcc' (Fold fun e acc)                      = combine "Fold" [ travFun f c l fun, travExp f c l  e, travAcc f c l acc]
      travAcc' (Fold1 fun acc)                       = combine "Fold1" [ travFun f c l fun, travAcc f c l acc]
      travAcc' (FoldSeg fun e acc1 acc2)             = combine "FoldSeg" [ travFun f c l fun, travExp f c l  e
                                                                         , travAcc f c l acc1, travAcc f c l acc2 ]
      travAcc' (Fold1Seg fun acc1 acc2)              = combine "FoldSeg1" [ travFun f c l fun, travAcc f c l acc1, travAcc f c l acc2 ]
      travAcc' (Scanl fun e acc)                     = combine "Scanl" [ travFun f c l fun, travExp f c l  e, travAcc f c l acc ]
      travAcc' (Scanl' fun e acc)                    = combine "Scanl'" [ travFun f c l fun, travExp f c l  e, travAcc f c l acc ]
      travAcc' (Scanl1 fun acc)                      = combine "Scanl1" [ travFun f c l fun, travAcc f c l acc ]
      travAcc' (Scanr fun e acc)                     = combine "Scanr" [ travFun f c l fun, travExp f c l  e, travAcc f c l acc ]
      travAcc' (Scanr' fun e acc)                    = combine "Scanr'" [ travFun f c l fun, travExp f c l  e, travAcc f c l acc ]
      travAcc' (Scanr1 fun acc)                      = combine "Scanr1" [ travFun f c l fun, travAcc f c l acc ]
      travAcc' (Permute fun dfts p acc)              = combine "Permute" [ travFun f c l fun, travAcc f c l dfts
                                                                         , travFun f c l p, travAcc f c l acc]
      travAcc' (Backpermute sh p acc)                = combine "Backpermute" [ travExp f c l  sh, travFun f c l p, travAcc f c l acc]
      travAcc' (Stencil sten bndy acc)               = combine "Stencil" [ travFun f c l sten, travBoundary f l acc bndy
                                                                         , travAcc f c l acc]
      travAcc' (Stencil2 sten bndy1 acc1 bndy2 acc2) = combine "Stencil2" [ travFun f c l sten, travBoundary f l acc1 bndy1
                                                                          , travAcc f c l acc1, travBoundary f l acc2 bndy2
                                                                          , travAcc f c l acc2]
      travAcc' (Collect seq)                         = travSeq f c l seq
  
  travSeq :: forall m b aenv senv a. Monad m => Labels -> (String -> String -> [m b] -> m b)
           -> (String -> String -> m b) -> PreOpenSeq OpenAcc aenv senv a -> m b
  travSeq f c l seq =
    case seq of
      Producer p s' ->
        case p of
          StreamIn _    -> combine "StreamIn" [ leaf "..." , travSeq f c l s' ]
          ToSeq slix _ a -> combine "ToSeq" [ leaf (travSlix slix), travAcc f c l a, travSeq f c l s' ]
          MapSeq afun x -> combine "MapSeq" [ travAfun f c l afun, leaf (show (idxToInt x)), travSeq f c l s' ]
          ChunkedMapSeq afun x -> combine "ChunkedMapSeq" [ travAfun f c l afun, leaf (show (idxToInt x)), travSeq f c l s' ]
          ZipWithSeq afun x y -> combine "ZipWithSeq" [ travAfun f c l afun, leaf (show (idxToInt x)), leaf (show (idxToInt y)), travSeq f c l s' ]
          ScanSeq fun e x -> combine "ScanSeq" [ travFun f c l fun, travExp f c l e, leaf (show (idxToInt x)), travSeq f c l s' ]
      Consumer co -> travC co
      Reify ix    -> leaf (show (idxToInt ix))
    where
      combine = c (accFormat f)
      leaf    = l (accFormat f)
      --
      travT :: Atuple (Consumer OpenAcc aenv senv) t' -> m b
      travT NilAtup          = l (tupleFormat f) "NilAtup"
      travT (SnocAtup tup s) = c (tupleFormat f) "SnocAtup" [ travT tup, travC s ]
  
      travC :: forall a. Consumer OpenAcc aenv senv a -> m b
      travC co =
        case co of
          FoldSeq fun e x -> combine "FoldSeq" [ travFun f c l fun, travExp f c l e, leaf (show (idxToInt x)) ]
          FoldSeqFlatten afun a x -> combine "FoldSeqFlatten" [ travAfun f c l afun, travAcc f c l a, leaf (show (idxToInt x)) ]
          Stuple t -> travT t
  
      travSlix :: SliceIndex slix sl co sh -> String
      travSlix SliceNil       = "Z"
      travSlix (SliceAll s)   = travSlix s ++ ":.All"
      travSlix (SliceFixed s) = travSlix s ++ ":.Split"
  
  travExp :: forall m env aenv a b . Monad m => Labels
         -> (String -> String -> [m b] -> m b)
         -> (String -> String -> m b)
         -> OpenExp env aenv a -> m b
  travExp f c l expr = travExp' expr
    where
      combine = c (expFormat f)
      leaf    = l (expFormat f)
  
      travExp' :: OpenExp env aenv a -> m b
      travExp' (Let e1 e2)                = combine "Let" [travExp f c l e1, travExp f c l e2]
      travExp' (Var idx)                  = leaf ("Var "   `cat` idxToInt idx)
      travExp' (Const v)                  = leaf ("Const " `cat` (toElt v :: a))
      travExp' (Tuple tup)                = combine "Tuple" [ travTuple f c l tup ]
      travExp' (Prj idx e)                = combine ("Prj " `cat` tupleIdxToInt idx) [ travExp f c l e ]
      travExp' (IndexNil)                 = leaf "IndexNil"
      travExp' (IndexCons t h)            = combine "IndexCons" [ travExp f c l t, travExp f c l h]
      travExp' (IndexHead ix)             = combine "IndexHead" [ travExp f c l ix ]
      travExp' (IndexTail ix)             = combine "IndexTail" [ travExp f c l ix ]
      travExp' (IndexAny)                 = leaf "IndexAny"
      travExp' (IndexSlice _ slix sh)     = combine "IndexSlice" [ travExp f c l slix, travExp f c l sh ]
      travExp' (IndexFull _ slix sl)      = combine "IndexFull" [ travExp f c l slix, travExp f c l sl ]
      travExp' (ToIndex sh ix)            = combine "ToIndex" [ travExp f c l sh, travExp f c l ix ]
      travExp' (FromIndex sh ix)          = combine "FromIndex" [ travExp f c l sh, travExp f c l ix ]
      travExp' (Cond cond thn els)        = combine "Cond" [travExp f c l cond, travExp f c l thn, travExp f c l els]
      travExp' (While cond body x)        = combine "While" [ travFun f c l cond, travFun f c l body, travExp f c l x ]
      travExp' (PrimConst a)              = leaf ("PrimConst " `cat` labelForConst a)
      travExp' (PrimApp p a)              = combine "PrimApp" [ l (primFunFormat f) (labelForPrimFun p), travExp f c l a ]
      travExp' (Index idx i)              = combine "Index" [ travAcc f c l idx, travExp f c l i]
      travExp' (LinearIndex idx i)        = combine "LinearIndex" [ travAcc f c l idx, travExp f c l i]
      travExp' (Shape idx)                = combine "Shape" [ travAcc f c l idx ]
      travExp' (ShapeSize e)              = combine "ShapeSize" [ travExp f c l e ]
      travExp' (Intersect sh1 sh2)        = combine "Intersect" [ travExp f c l sh1, travExp f c l sh2 ]
      travExp' (Union sh1 sh2)            = combine "Union" [ travExp f c l sh1, travExp f c l sh2 ]
      travExp' (Foreign ff fun e)         = combine ("Foreign " ++ strForeign ff) [ travFun f c l fun, travExp f c l e ]
  
  
  travAfun :: forall m b aenv fun. Monad m => Labels -> (String -> String -> [m b] -> m b)
          -> (String -> String -> m b) -> OpenAfun aenv fun -> m b
  travAfun f c l openAfun = travAfun' openAfun
    where
      combine = c (funFormat f)
      travAfun' :: OpenAfun aenv fun -> m b
      travAfun' (Abody body) = combine "Abody" [ travAcc f c l body ]
      travAfun' (Alam fun)   = combine "Alam"  [ travAfun f c l fun ]
  
  travArrays
      :: forall m a b. Monad m
      => Labels
      -> (String -> String -> [m b] -> m b)
      -> (String -> String -> m b)
      -> ArraysR a
      -> a
      -> m b
  travArrays f c l = trav
    where
      combine = c (accFormat f)
      leaf    = l (accFormat f)
      --
      trav :: ArraysR a' -> a' -> m b
      trav ArraysRunit         ()       = leaf    "ArraysRunit"
      trav ArraysRarray        a        = combine "ArraysRarray" [ travArray f l a ]
      trav (ArraysRpair r1 r2) (a1, a2) = combine "ArraysRpair"  [trav r1 a1, trav r2 a2]
  
  travArray :: forall dim a m b. Monad m => Labels -> (String -> String -> m b) -> Array dim a -> m b
  travArray f l (Array sh _) = l (arrayFormat f) ("Array" `cat` (toElt sh :: dim))
  
  travBoundary :: forall aenv dim e m b. (Monad m, Elt e) => Labels -> (String -> String -> m b)
               -> {-dummy-}OpenAcc aenv (Array dim e)
               -> Boundary (EltRepr e) -> m b
  travBoundary f l boundary = travBoundary' boundary
    where
      leaf = l (boundaryFormat f)
      travBoundary' :: {-dummy-}OpenAcc aenv (Array dim e) -> Boundary (EltRepr e) -> m b
      travBoundary' _ Clamp        = leaf "Clamp"
      travBoundary' _ Mirror       = leaf "Mirror"
      travBoundary' _ Wrap         = leaf "Wrap"
      travBoundary' _ (Constant e) = leaf ("Constant " `cat` (toElt e :: e))
  
  
  travAtuple
      :: forall m b aenv t. Monad m
      => Labels
      -> (String -> String -> [m b] -> m b)
      -> (String -> String -> m b)
      -> Atuple (OpenAcc aenv) t
      -> m b
  travAtuple f c l = trav
    where
      leaf    = l (tupleFormat f)
      combine = c (tupleFormat f)
      --
      trav :: Atuple (OpenAcc aenv) t' -> m b
      trav NilAtup          = leaf    "NilAtup"
      trav (SnocAtup tup a) = combine "SnocAtup" [ trav tup, travAcc f c l a ]
  
  
  travFun :: forall m b env aenv fun.Monad m => Labels -> (String -> String -> [m b] -> m b)
          -> (String -> String -> m b) -> OpenFun env aenv fun -> m b
  travFun f c l openFun = travFun' openFun
    where
      combine = c (funFormat f)
      travFun' :: OpenFun env aenv fun -> m b
      travFun' (Body body) = combine "Body" [ travExp f c l body ]
      travFun' (Lam fun)   = combine "Lam"  [ travFun f c l fun ]
  
  
  travTuple :: forall m b env aenv t. Monad m => Labels -> (String -> String -> [m b] -> m b)
            -> (String -> String -> m b) -> Tuple (OpenExp env aenv) t -> m b
  travTuple f c l tuple = travTuple' tuple
    where
      leaf    = l (tupleFormat f)
      combine = c (tupleFormat f)
      travTuple' :: Tuple (OpenExp env aenv) t -> m b
      travTuple' NilTup          = leaf     "NilTup"
      travTuple' (SnocTup tup e) = combine "SnocTup" [ travTuple f c l tup, travExp f c l e ]
  
  labelForPrimFun :: PrimFun a -> String
  labelForPrimFun (PrimAdd _)            = "PrimAdd"
  labelForPrimFun (PrimSub _)            = "PrimSub"
  labelForPrimFun (PrimMul _)            = "PrimMul"
  labelForPrimFun (PrimNeg _)            = "PrimNeg"
  labelForPrimFun (PrimAbs _)            = "PrimAbs"
  labelForPrimFun (PrimSig _)            = "PrimSig"
  labelForPrimFun (PrimQuot _)           = "PrimQuot"
  labelForPrimFun (PrimRem _)            = "PrimRem"
  labelForPrimFun (PrimQuotRem _)        = "PrimQuotRem"
  labelForPrimFun (PrimIDiv _)           = "PrimIDiv"
  labelForPrimFun (PrimMod _)            = "PrimMod"
  labelForPrimFun (PrimDivMod _)         = "PrimDivMod"
  labelForPrimFun (PrimBAnd _)           = "PrimBAnd"
  labelForPrimFun (PrimBOr _)            = "PrimBOr"
  labelForPrimFun (PrimBXor _)           = "PrimBXor"
  labelForPrimFun (PrimBNot _)           = "PrimBNot"
  labelForPrimFun (PrimBShiftL _)        = "PrimBShiftL"
  labelForPrimFun (PrimBShiftR _)        = "PrimBShiftR"
  labelForPrimFun (PrimBRotateL _)       = "PrimBRotateL"
  labelForPrimFun (PrimBRotateR _)       = "PrimBRotateR"
  labelForPrimFun (PrimFDiv _)           = "PrimFDiv"
  labelForPrimFun (PrimRecip _)          = "PrimRecip"
  labelForPrimFun (PrimSin _)            = "PrimSin"
  labelForPrimFun (PrimCos _)            = "PrimCos"
  labelForPrimFun (PrimTan _)            = "PrimTan"
  labelForPrimFun (PrimAsin _)           = "PrimAsin"
  labelForPrimFun (PrimAcos _)           = "PrimAcos"
  labelForPrimFun (PrimAtan _)           = "PrimAtan"
  labelForPrimFun (PrimSinh _)           = "PrimSinh"
  labelForPrimFun (PrimCosh _)           = "PrimCosh"
  labelForPrimFun (PrimTanh _)           = "PrimTanh"
  labelForPrimFun (PrimAsinh _)          = "PrimAsinh"
  labelForPrimFun (PrimAcosh _)          = "PrimAcosh"
  labelForPrimFun (PrimAtanh _)          = "PrimAtanh"
  labelForPrimFun (PrimExpFloating _)    = "PrimExpFloating"
  labelForPrimFun (PrimSqrt _)           = "PrimSqrt"
  labelForPrimFun (PrimLog _)            = "PrimLog"
  labelForPrimFun (PrimFPow _)           = "PrimFPow"
  labelForPrimFun (PrimLogBase _)        = "PrimLogBase"
  labelForPrimFun (PrimTruncate _ _)     = "PrimTruncate"
  labelForPrimFun (PrimRound _ _)        = "PrimRound"
  labelForPrimFun (PrimFloor _ _)        = "PrimFloor"
  labelForPrimFun (PrimCeiling _ _)      = "PrimCeiling"
  labelForPrimFun (PrimAtan2 _)          = "PrimAtan2"
  labelForPrimFun (PrimIsNaN _)          = "PrimIsNaN"
  labelForPrimFun (PrimLt _)             = "PrimLt"
  labelForPrimFun (PrimGt _)             = "PrimGt"
  labelForPrimFun (PrimLtEq _)           = "PrimLtEq"
  labelForPrimFun (PrimGtEq _)           = "PrimGtEq"
  labelForPrimFun (PrimEq _)             = "PrimEq"
  labelForPrimFun (PrimNEq _)            = "PrimNEq"
  labelForPrimFun (PrimMax _)            = "PrimMax"
  labelForPrimFun (PrimMin _)            = "PrimMin"
  labelForPrimFun PrimLAnd               = "PrimLAnd"
  labelForPrimFun PrimLOr                = "PrimLOr"
  labelForPrimFun PrimLNot               = "PrimLNot"
  labelForPrimFun PrimOrd                = "PrimOrd"
  labelForPrimFun PrimChr                = "PrimChr"
  labelForPrimFun PrimBoolToInt          = "PrimBoolToInt"
  labelForPrimFun (PrimFromIntegral _ _) = "PrimFromIntegral"
  
  labelForConst :: PrimConst a -> String
  labelForConst (PrimMinBound _) = "PrimMinBound"
  labelForConst (PrimMaxBound _) = "PrimMaxBound"
  labelForConst (PrimPi       _) = "PrimPi"
  
  -- Auxiliary ops
  --
  
  