packages feed

fregel-1.2.0: compiler/Normalization.hs

{-# Language TypeSynonymInstances,FlexibleInstances,MultiParamTypeClasses,FunctionalDependencies,RankNTypes,FlexibleContexts,KindSignatures,ScopedTypeVariables #-}
module Normalization where

import Control.Monad.State
import Data.Maybe
import Data.List
import Numeric (showHex)
import Debug.Trace

import Spec
import ASTData
import Analysis (dependencyC, topologicalSort)
import TypeChecker (typing, unify, apply, buildVertCompEnv, buildProgEnv, typing2', typeCheck)
import Inlining (runInlining)
import AggregatorExtraction (aggExtraction, aggExtractionE)
import TypeInstantiation (runTypeInstantiation)
import DependencySimple (computeDependency)
import VCOutputNormalization (normalizeVertCompRecordOutputExp, normalizeVertCompRecordOutputExp')

{-


Big-step normalization. See ICFP paper.

We can make a small-step normalization (by fusion/tupling/...), but leads to inefficient code (or difficult-to-be-optimized code).

This module drives all compilation steps except for code generation: normalizationAll .

The interface to code generators is the data structure named DNormalized.

-}
{-

**********************
Normalized Fregel Code
**********************
A single 'fregel' computation that executes a vertex-computation function (ft_Pi) according to the current 'phase' (i).

 n_fregel NData initPhase [(ft_P1), ..., (ft_Pk)] next stay g

The parameters (types and vertex-computation functions):

  data Phase = P1 | P2 | ... | Pk  # unique names of the phases
  data NData = NData { phase::Maybe Phase,             -- the current phase
                       step::Int,                      -- the number of steps in the current phase
                       data_P1::Pair Bool Data_Pi,     -- the pair of end-flag and data for phase P1
                       ...,                            -- and so on
                       data_Pk::Pair Bool Data_Pi,
                       step_X1::Int,                   -- the number of iterations in the first giter
                       ...,                            -- and so on
                       step_Xl::Int }
  ft_Pi v prev curr :: Pair Bool Data_Pi               -- vertex-computation function for phase Pi
                                                           built by composing the initialization func.,
                                                           step funct. and termination judgement func.
                                                           of the phase.
  initPhase :: Phase                                   -- the initial phase
  next :: Phase -> Phase        -- returns the phase to be executed after the given phase terminates 
  stay :: Phase -> Phase        -- returns the phase to be executed when the given phase continues 
                                -- stay p = p for most phases,
                                --  but for giter phases stay p is the first phase of its internal ones

Implementation:
It executes a vertex-computation function (ft_Pi) according to the current 'phase' (i).
If the termination condition for the current phase is satisfied, then it proceeds to the next phase.

n_fregel NData initPhase [(ft_P1), ..., (ft_Pk)] next stay g
 = let ft v prev curr =
         let d1 = if prev v.^phase == Just P1 then ft_P1 v prev curr else prev v.^data_P1
             ...
             dk = if prev v.^phase == Just Pk then ft_Pk v prev curr else prev v.^data_Pk
             p_end  = if prev v.^phase == Just P1 then curr v.^data_P1.^_fst
                      else if prev v.^phase == Just P2 then curr v.^data_P2.^_fst
                           else ...
                                else if prev v.^phase == Just Pk then curr v.^data_Pk.^_fst
                                     else False
             step' = if6 not (prev v.^phase == Nothing) && prev v.^phase == curr v.^phase then prev v.^step + 1 else 0
             phase' = if p_end then next (prev v.^phase) else stay (prev v.^phase)
             step_X1' = if prev v.^phase == Just X1 then if not p_end then prev v.^step_X1 + 1 else 0 else prev v.^step_X1
             ...
             step_Xl' = if prev v.^phase == Just Xl then if not p_end then prev v.^step_Xl + 1 else 0 else prev v.^step_Xl
         in NData phase' step' d1 ... dk step_X1' ... step_Xl'
       f0 v = NData (Just initPhase) 0 _ ... _
   in fregel f0 ft Fix g

-}


{---  data structure for Normalized Code (i.e., phase-transition machine)  ---}

type DPhaseID = Int
type DPhase a = (DPhaseID, String, DDefVertComp a, DDefVertComp a, DDefVertComp a)
   -- Each phase is represented by ID, label, initialization function, step function, and terminatino judgement function
type DPhaseGraph a = [(DPhaseID, [([DField a], DPhaseID)])]  -- this is currently used to represent the stay and next functions: [ (p, [([], stay p), ([], next p)]) | p in all_phases ]. In general, this was indended to represent dependencies among phases and the list [DField a] was used to remember the ways of dependencies.

data DNormalized a = DNormalized
                         String          -- the original program (function) name
                         [DPhaseID]      -- giter phases (these will have their own iteration counters)
                         [DRecordSpec a] -- user-defined records
                         (DRecordSpec a, [DPhaseID]) -- NData (vertex data + phase info.) and corresponding phase IDs
                         --TODO: (DRecordSpec a, [DPhaseID], DPhaseID) -- NData (vertex data + phase info.) and corresponding phase IDs and the phase of the final result (>=0 or -1 for N.A.)
                         [DPhase a]      -- all phases
                         (DPhaseGraph a) -- dependencies among phases (currently, stay and next funcs.)
                         [DPhaseID]      -- the initial phases (currently, only one)
                         [DSmplDef a]    -- global definitions (such as constants)
                         deriving (Eq, Show)

-- the big-step normalization to build data representing the n_fregel program.           
makeNormalizedData :: (DProgramSpec DASTData, DUnique) -> DNormalized DASTData
makeNormalizedData (p, uid) = 
    let s = runExtractGV (p, uid)
        (DProgramSpec rs (DProg f defs _ _) a) = p
        fn = getName f
        daNew = DASTData (newDataType fn) []
        c = (DConstructor (newDataTypeName fn) daNew)
        daInt = (DASTData typeDTInt [])
        phaseSystemData = [(DField (fieldPhase fn) daInt, DTInt daInt), (DField (fieldStep fn) daInt, DTInt daInt)]
        iterCounters = map (\n -> (DField (mkStep n) daInt, DTInt daInt)) (iters s)
        sortWithCurr' = sortWithCurr (map extractFields rs)
        mkPhase (i::Int, (n, t, _, _, (f1, f2, f3))) = 
            let phaseDataFieldName = mkField n
            in (i, n, sortWithCurr' [phaseDataFieldName, "_snd"] f1, sortWithCurr' [phaseDataFieldName, "_snd"] f2, sortWithCurr' [phaseDataFieldName, "_fst"] f3)
        graphVarsWithIds = zip [0..] (reverse $ head $ vars s) -- this assigns phase IDs
        phases = map mkPhase graphVarsWithIds
        fsWithIds = map (\(pid, (n, t, _, _, _)) -> let da = (DASTData t []) in ((DField (mkField n) da, buildTypeAST (typePair [typeDTBool, t])), pid)) graphVarsWithIds
        newDataFiledsWithIds = map (\x -> (x, -1)) phaseSystemData++fsWithIds++map (\x -> (x, -1)) iterCounters
        fs = map fst newDataFiledsWithIds -- fields of NData
        ps = map snd newDataFiledsWithIds -- phase ID of each field in NData
        newData = DRecordSpec c fs daNew
        idmap = map (\(i::Int, n::DVarName, _, _, _) -> (n, i)) phases
        lookupId nm = let (Just i) = lookup nm idmap in i
        initPhase = lookupId (maybe (error "something wrong") (\a -> a) $ initVar s)
        defs' = map (\(DGDefSmpl def _) -> def) $ filter (isSmplDef) defs
        deps = map (\(v, _, n, s, _) -> (lookupId v, [([], lookupId s)] ++ if n == "" then [] else [([], lookupId n)])) (reverse $ head $ vars s) 
        is = map (\g -> lookupId g) $ iters s
    in DNormalized fn is rs (newData, ps) phases deps [initPhase] defs'

-- misc functions for the big-step normalizer
       
extractFields :: DRecordSpec a -> (String, [String])
extractFields (DRecordSpec (DConstructor c _) fts _) = (c, map (\(DField f _, _) -> f) fts)

buildTypeAST (t@(DTypeTerm "Int" [])) = DTInt (DASTData t [])
buildTypeAST (t@(DTypeTerm "Bool" [])) = DTBool (DASTData t [])
buildTypeAST (t@(DTypeTerm "String" [])) = DTString (DASTData t [])
buildTypeAST (t@(DTypeTerm "Double" [])) = DTDouble (DASTData t [])
buildTypeAST (t@(DTypeTerm "(,)" ts)) = DTTuple (map buildTypeAST ts) (DASTData t [])
buildTypeAST (t@(DTypeTerm c ts)) = DTRecord (DConstructor c (DASTData t [])) (map buildTypeAST ts) (DASTData t [])

buildTypeAST (t) = error ("Unknown situation in buildTypeAST: " ++ show t)

isSmplDef (DGDefSmpl def a) = True
isSmplDef _ = False

newDataTypeName fn = "NData_"++ fn
fieldPhase fn = "phase_"++fn
fieldStep fn = "step_"++fn
newDataType fn = typeSolid (newDataTypeName fn)
mkField gx = "data_"++gx
mkStep gx = "step_"++gx

                                                                             
findInputGraphType (DProgramSpec _ p _) = let (DTypeTerm _ (igt:_)) = typeOf (getData p)
                                          in igt

getRecordsSpecs = do s <- get
                     return (recordsSpecs s)
                                             
-- extracting graph variables to collect information to build the phase-transition machine
runExtractGV ((p@(DProgramSpec rs (DProg f defs e af) a)), uid)
    = let igt = (findInputGraphType p)
          (_, s) = runState (extractGV p) (DEnvGV [] [] [] Nothing defs igt [] (getName f) uid rs)
      in s { groundDefs = [] }

-- the state: the data collected during the extraction
data DEnvGV = DEnvGV { vars :: [Chain],                   -- a stack of chains under extracting
                       iters :: [DVarName],               -- giter labels
                       loops :: [(DVarName, Chain)],      -- a graph function = a chain
                       initVar :: Maybe DVarName,         -- the entry point (initial phase)
                       groundDefs::[DGroundDef DASTData], -- static: definitions of graph functions (used to lookup their input graph names)
                       inputGraphType::DTypeInfo,         -- static: the input graph type (of the whole computation)
                       typeInfoOf :: [(DVarName, DTypeInfo)], -- types of labels found so far 
                       theProgName :: DVarName,           -- static: the program name
                       uniqId :: Int,                     -- uniq id...
                       recordsSpecs :: [DRecordSpec DASTData] -- record specs.
                     } deriving (Eq)

instance (Show DEnvGV) where
    show (DEnvGV vs is _ iv _ igt tis fn uid rs) = unlines (["Program: " ++ fn] ++ vss++[iterss, ivs, igts, tiss])
        where iterss = "iters = " ++ show is
              ivs = "initVar = " ++ show iv
              tiss = "Types = " ++ prettyShow tis ++ " --EndOfTypes"
              igts = "inputGraphType = " ++ show igt
              vss = concatMap showVars (concat vs)
              showVars (v, t, n, s, (fi, ft, fp)) = ["Phase: " ++ v, "Type: " ++ prettyShow t, "Next: " ++ n, "Stay: " ++ s] ++ ppDefVertComp fi ++ ppDefVertComp ft ++ ppDefVertComp fp


-- the triplet for a phase: initilization, step and judgement functions
type PhaseFunctions = (DDefVertComp DASTData, DDefVertComp DASTData, DDefVertComp DASTData)

-- list of (phase label (= graph var. name), type, next, stay, the triplet)
-- a chain = a graph function
type Chain = [(DVarName, DTypeInfo, DVarName, DVarName, PhaseFunctions)]
nameOf (v, t, n, s, fs) = v

-- accessors to the state
setProgName fn = do env <- get
                    put (env {theProgName = fn})

getProgName = do env <- get
                 return (theProgName env)

getUid :: State DEnvGV DUnique
getUid = do env <- get
            return (uniqId env)

setUid uid = do env <- get
                put (env { uniqId = uid })

getInputGraphType :: State DEnvGV DTypeInfo
getInputGraphType = do env <- get
                       return (inputGraphType env)

-- update the next phase of the head of the given chain to the given var
updateHead v' [] = []
updateHead v' ((v, t, n, s, fs):xs) = (v, t, v', s, fs):xs

getDefs :: State DEnvGV [DGroundDef DASTData]
getDefs = do env <- get
             return (groundDefs env)

addTypeInfo v t = 
    do env <- get
       put (env {typeInfoOf = (v, t):(typeInfoOf env)})

-- adds a new var to the current chain
addVar v = addVar' v (getName v)

addVar' :: DVar DASTData -> String -> PhaseFunctions -> State DEnvGV ()
addVar' v stay fs = 
    do env <- get
       let vs = head $ vars env
           vss = tail $ vars env
           types = typeInfoOf env
           label = getName v
           theType = getVertexType (typeOf (getData v))
           env' = env { vars = (([(label, theType, "", stay, fs)] ++ updateHead (getName v) vs) : vss), typeInfoOf = (label, theType) :types }
       if length (vars env') == 1 && initVar env' == Nothing then put (env' { initVar = Just (getName v) } ) else put env'

-- starts a new chain of graph vars
newVarChain :: State DEnvGV ()
newVarChain = do env <- get
                 put ( env { vars = []:vars env } )

-- pops the current chain
popVarChain :: State DEnvGV Chain
popVarChain = do env <- get
                 put (env { vars = tail (vars env) })
                 return (head $ vars env)

-- save a chain to be used later to make a loop by iter
saveChain :: DFun DASTData -> Chain -> State DEnvGV ()
saveChain f chain = do env <- get
                       put (env { loops = (getName f, chain) : loops env })

-- remove the chain indexed by f
removeChain :: DFun DASTData -> State DEnvGV ()
removeChain f = do env <- get
                   put (env { loops = filter (\(n, _) -> not (n == getName f)) (loops env)})

-- find the chain indexed by f
findChain :: DFun DASTData -> State DEnvGV Chain
findChain f = do env <- get
                 case lookup (getName f) (loops env) of
                   Nothing -> error "something wrong with dependency info.?"
                   Just chain -> do removeChain f --single use
                                    return chain
-- combination of find/remove
popChain :: DFun DASTData -> State DEnvGV Chain
popChain f = do chain <- findChain f
                removeChain f
                return chain

-- replacing the argument graph name of a graph function (a chain) with the actual input graph name
replaceInputGraphName gi gx (v, t, n, s, (f0, ft, fj)) 
    = let f0' = repVC f0
          ft' = repVC ft
          fj' = repVC fj
          --TODO: update dep
          repVC (DDefVertComp f defs e a) = 
              let defs' = map repD defs
                  e' = rec e
              in DDefVertComp f defs' e' a
          repD (DDefVar v [] e a) = 
              let e' = rec e in DDefVar v [] e' a
          repD (DDefTuple vs [] e a) = 
              let e' = rec e in DDefTuple vs [] e' a
      in (v, t, n, s, (f0', ft', fj')) 
  where
    rec (DIf p t e a) = 
        let p' = rec p
            t' = rec t
            e' = rec e
        in (DIf p' t' e' a)
    rec (DTuple es a) =
        let es' = map rec es 
        in (DTuple es' a)
    rec (DFunAp f es a) =
        let es' = map rec es
        in (DFunAp f es' a)
    rec (DConsAp c es a) =
        let es' = map rec es
        in (DConsAp c es' a)
    rec ((DFieldAcc t fs a)) = 
        let fs' = map repF fs
        in DFieldAcc t fs' a
    rec (x@(DFieldAccE e fs a)) = x
    rec (DAggr a' e g es a) =
        let e'  = rec e
            es' = map rec es
        in (DAggr a' e' g es' a)
    rec (x@(DVExp (DVar v av) a)) = x
    rec (x@(DCExp c a)) = x
    repF (DField f a) = DField (if f == f_gi then f_gx else f) a
    f_gi = mkField gi
    f_gx = mkField gx
    
-- for iter-label v, closing the chain of f and adding the label as the last in the current chain
addClosedChain v f e types = 
    do chain <- popChain f
       defs <- getDefs
       let chain' = updateHead (getName v) chain
           (Just (DGDefGF (DDefGraphFun _ gi _ _ _) _)) = lookupBy getName (getName f) defs
           chain'' = map (replaceInputGraphName (getName gi) (getName v)) chain
       env <- get
       let vs = head $ vars env
           vss = tail $ vars env
           startOfChain = nameOf $ last chain''
           env' = env { vars = (chain'' ++ updateHead (getName v) vs) : vss } -- add the chain
       put env'
       fn <- getProgName
       fs <- (buildFunctions fn e v types defs)
       addVar' v startOfChain fs

-- adds an iter label
addIter v = do env <- get
               let env' = env { iters = getName v : iters env }
               put env'

getTypeInfo :: State DEnvGV [(DVarName, DTypeInfo)]
getTypeInfo = do env <- get
                 return (typeInfoOf env)

-- visitor to extract graph variables (build chains) 
class GraphVarExtractable a where
    extractGV :: a -> State DEnvGV ()

instance GraphVarExtractable (DProgramSpec DASTData) where
    extractGV (DProgramSpec rs p a) = extractGV p

instance GraphVarExtractable (DProg DASTData) where
    extractGV (DProg f defs e a) = 
        do newVarChain
           mapM_ extractGV defs

instance GraphVarExtractable (DGroundDef DASTData) where
    extractGV (DGDefVC d a) = return ()
    extractGV (DGDefVI d a) = return ()
    extractGV (DGDefSmpl d a) = return ()
    extractGV (DGDefGV d a) = extractGV d
    extractGV (DGDefGF d a) = extractGV d

instance GraphVarExtractable (DDefGraphVar DASTData) where
    extractGV (DDefGraphVar v e a) = 
        do types' <- getTypeInfo
           let types = types' ++ [(getName v, getVertexType (typeOf a))]
           --trace (getName v ++ "is being processed with types = "++show types) $
           case e of 
             (DGIter f0 ft x g _) -> do addClosedChain v ft e types
                                        addIter v
             otherwise            -> do defs <- getDefs
                                        fn <- getProgName
                                        fs <- (buildFunctions fn e v types defs)
                                        addVar v fs

------------------- substitutions to build the triplet ------------------                      
-- for ft/f0
sigma fn go gi types typeVofG typeEofG = recToFieldAcc replace 
  where
    replace (DFieldAcc (x@(DPrev v av)) fs a) = 
        let (av', fs') = addField go fs a av types fn
            v' = setData ((getData v) { typeOf = typeGraph [typeVofG, typeEofG]}) v
        in DFieldAcc (DPrev v' av') fs' a
    replace (DFieldAcc (x@(DCurr v av)) fs a) = 
        let (av', fs') = addField go fs a av types fn
            v' = setData ((getData v) { typeOf = typeGraph [typeVofG, typeEofG]}) v
        in DFieldAcc (DCurr v' av') fs' a
    replace (y@(DFieldAcc (x@(DVal v av)) fs a)) =
        if gi == "g" then y else let (av', fs') = addField gi fs a av types fn
                                 in DFieldAcc (DPrev v av') fs' a

-- for ft/f0 with gzip-elimination optimization: field accessors will be replaced.
sigmaZ gez fn go types typeVofG typeEofG = recToFieldAcc replace 
  where
    replace (DFieldAcc (x@(DPrev v av)) fs a) = 
        let (av', fs') = addField go fs a av types fn
            v' = setData ((getData v) { typeOf = typeGraph [typeVofG, typeEofG]}) v
        in DFieldAcc (DPrev v' av') fs' a
    replace (DFieldAcc (x@(DCurr v av)) fs a) = 
        let (av', fs') = addField go fs a av types fn
            v' = setData ((getData v) { typeOf = typeGraph [typeVofG, typeEofG]}) v
        in DFieldAcc (DCurr v' av') fs' a
    replace (y@(DFieldAcc (x@(DVal v av)) fs a)) =
        case isIncompleteAccess fs gez of
          Nothing -> addFieldZ gez fs a av types fn v
          Just rest -> buildPairExpression v rest -- needs to build a pair locally here.                                 
    buildPairExpression v (DGVar gv a)
        = let tv = getVertexType (typeOf a)
              dummy = DASTData tv []
              (av', fs') = addField (getName gv) [] dummy dummy types fn
          in DFieldAcc (DPrev v av') fs' (DASTData tv []) -- TODO: add dependency data...
    buildPairExpression v (DGZip ge1 ge2 a)
        = let pe1 = buildPairExpression v ge1
              pe2 = buildPairExpression v ge2
              t1 = getVertexType (getType ge1)
              t2 = getVertexType (getType ge2)
              pt = typePair [t1, t2]
              a' = DASTData pt [] -- TODO: add dependency data...
              ac = DASTData (typeFunction [t1, t2, pt]) [] 
          in expConstructor' "Pair" [pe1, pe2] ac a'
                                     
isIncompleteAccess [] (gez@(DGZip _ _ _)) = Just gez  --needs to build a pair locally!
isIncompleteAccess [] _ = Nothing
isIncompleteAccess (f:fs) (ge@(DGZip ge1 ge2 a))
    = case f of DField "_fst" _ -> isIncompleteAccess fs ge1 
                DField "_snd" _ -> isIncompleteAccess fs ge2
                otherwise -> error ("why " ++ show f ++ " to access " ++ show ge)
isIncompleteAccess (f:fs) _ = Nothing

-- for Until (\g- > e)
sigma' fn gx types typeVofG typeEofG = recToFieldAcc replace 
  where
    replace (y@(DFieldAcc (x@(DVal v av)) fs a)) = 
        let (av', fs') = addField gx fs a av types fn
        in DFieldAcc (DCurr v av') fs' a

addField gx fs a av types fn = (av', fs')
        where t2 = typePair [typeDTBool, t1]
              t1 = case lookup gx types of (Just x) -> x; Nothing -> error ("why? " ++ gx ++ " is not found in " ++ show types)
              a2 = a { typeOf = t2 }
              a1 = a { typeOf = t1 }
              fs' = [ DField (mkField gx) a2 , DField "_snd" a1 ] ++ fs
              av' = av { typeOf = newDataType fn }

addFieldZ gez fs0 a av types fn v = DFieldAcc (DPrev v av') fs' a
        where (fs, gx) = downToGVar fs0 gez
              t2 = typePair [typeDTBool, t1]
              t1 = case lookup gx types of (Just x) -> x; Nothing -> error ("why? " ++ gx ++ " is not found in " ++ show types)
              a2 = a { typeOf = t2 }
              a1 = a { typeOf = t1 }
              fs' = [ DField (mkField gx) a2 , DField "_snd" a1 ] ++ fs
              av' = av { typeOf = newDataType fn }

downToGVar fs (DGVar gx a) = (fs, getName gx)
downToGVar (f:fs) (DGZip ge1 ge2 a)
    = case f of DField "_fst" _ -> downToGVar fs ge1 
                DField "_snd" _ -> downToGVar fs ge2

                    
recToFieldAcc replace = rec
  where
    rec (DIf p t e a) = 
        let p' = rec p
            t' = rec t
            e' = rec e
        in (DIf p' t' e' a)
    rec (DTuple es a) =
        let es' = map rec es 
        in (DTuple es' a)
    rec (DFunAp f es a) =
        let es' = map rec es
        in (DFunAp f es' a)
    rec (DConsAp c es a) =
        let es' = map rec es
        in (DConsAp c es' a)
    rec (x@(DFieldAcc t fs a)) = replace x
    rec (x@(DFieldAccE e fs a)) = x
    rec (DAggr a' e g es a) =
        let e'  = rec e
            es' = map rec es
        in (DAggr a' e' g es' a)
    rec (x@(DVExp v a)) = x
    rec (x@(DCExp c a)) = x

applySigma s (DDefVertComp f defs e a) typeVofG typeEofG = 
    let defs' = map (applySigmaD s) defs
        e' = s e
        a' = a    -- TODO: type has to be changed because of the change of curr/prev's types
        f' = setData a' f
    in (DDefVertComp f' defs' e' a')

applySigmaD s (x@(DDefVar v [] e a)) = DDefVar v [] (s e) a  -- the type is kept during the rewriting
applySigmaD s (x@(DDefTuple vs [] e a)) = DDefTuple vs [] (s e) a -- the type is kept during the rewriting

delDefs (DDefVertComp f _ e a) = DDefVertComp f [] e a

-- TODO: correct VertComp' type according to the output type of prev/curr?

-------- the triplet builders ----------------------

-- building (the initial function, the step function, the judgement function) for a given graph exrepssion
buildFunctions :: DVarName -> DGraphExpr DASTData -> DVar DASTData -> [(DVarName, DTypeInfo)] -> [DGroundDef DASTData] -> State DEnvGV PhaseFunctions
-- assumption: the argument is GDVar
buildFunctions fn (ge@(DPregel f0 ft x (DGVar g _) af)) v types defs =
    case x of 
      (DTermF _)   -> buildFunctionsPregel buildDefsBodyFp fn f0 ft g v types defs
      (DTermI e _) -> buildFunctionsPregel (buildDefsBodyIp e) fn f0 ft g v types defs
      (DTermU e _) -> buildFunctionsPregel (buildDefsBodyUp e) fn f0 ft g v types defs
                      
-- assumption: the argument is GDVar
buildFunctions fn (ge@(DGIter f0 ft x (DGVar g _) af)) v types defs =
    case x of 
      (DTermF _)   -> buildFunctionsIter buildDefsBodyFi fn f0 ft g v types defs
      (DTermI e _) -> buildFunctionsIter (buildDefsBodyIi e) fn f0 ft g v types defs 
      (DTermU e _) -> buildFunctionsIter (buildDefsBodyUi e) fn f0 ft g v types defs 
                      
-- assumption: the argument is GDVar
buildFunctions fn (ge@(DGMap f0 (DGVar g _) af)) v types defs =
    do gt <- getInputGraphType
       let typeVofG = getVertexType gt
           typeEofG = getEdgeType gt
           label = getName v
           theType = lookupType label types
           s = sigma fn label (getName g) types typeVofG typeEofG
           f0' = let (Just (DGDefVI (DDefVertInit f defs' e af) a)) = lookupBy getName (getName f0) defs in DDefVertComp f defs' e af
           body = expBool True -- phase-termination = always true
       xx <- typingForBuildFunctions typeVofG typeEofG theType fn [body] [] []
       let [body'] = xx
           [fj'] = buildDefVCtyped typeVofG typeEofG theType label [("_judge", typeDTBool, [], body')]   
       f0'' <- normalizeVertCompRecordOutputExp'' f0'
       return (applySigma s f0'' typeVofG typeEofG, applySigma s (delDefs f0'') typeVofG typeEofG {- this is dummy data, removing inner defs to avoid name conflicts -}, fj')

-- assumption: the argument is nested gzips!
buildFunctions fn (ge@(DGMap f0 gez af)) v types defs =
    do gt <- getInputGraphType
       let typeVofG = getVertexType gt
           typeEofG = getEdgeType gt
           label = getName v
           theType = lookupType label types
               -- here, we need better sigma that replaces (_fst|_snd)* with suitable counterparts in gez
           s = sigmaZ gez fn label types typeVofG typeEofG
           f0' = let (Just (DGDefVI (DDefVertInit f defs' e af) a)) = lookupBy getName (getName f0) defs in DDefVertComp f defs' e af
           body = expBool True -- phase-termination = always true
       xx <- typingForBuildFunctions typeVofG typeEofG theType fn [body] [] []
       let [body'] = xx
           [fj'] = buildDefVCtyped typeVofG typeEofG theType label [("_judge", typeDTBool, [], body')]   
       f0'' <- normalizeVertCompRecordOutputExp'' f0'
       return (applySigma s f0'' typeVofG typeEofG, applySigma s (delDefs f0'') typeVofG typeEofG {- this is dummy data, removing inner defs to avoid name conflicts -}, fj')

-- TODO: implement nested zips... Hmm, the code generator does someting special to gzip case..... pending.
              
-- assumption: the arguments are GDVar
buildFunctions fn (ge@(DGZip (DGVar g1 _) (DGVar g2 _) af)) v types defs =
    do gt <- getInputGraphType
       let typeVofG = getVertexType gt
           typeEofG = getEdgeType gt
           label = getName v
           label1 = getName g1               
           label2 = getName g2
           theType = lookupType label types
           typeGV1 = lookupType label1 types
           typeGV2 = lookupType label2 types
           -- the function body (pairing the values from g1 and g2)
           f0body = expConstructor "Pair" [(expFieldAcc "prev" "v" [mkField label1,"_snd"]), (expFieldAcc "prev" "v" [mkField label2,"_snd"])]
           body = expBool True -- phase-termination = always true
       xx <- {- trace (prettyShow typeG1 ++ "," ++ prettyShow typeG2) $ -} typingForBuildFunctions typeVofG typeEofG theType fn [f0body, body] [(label, theType), (label1, typeGV1), (label2, typeGV2)] []
       let [f0body', body'] = xx
           [fj', f0'] = buildDefVCtyped typeVofG typeEofG theType label [("_judge", typeDTBool, [], body'), ("_f0", theType, [], f0body')]   
       f0'' <- normalizeVertCompRecordOutputExp'' f0'
       -- Hmm.. using f0'' here causes an error in NtoPregel.hs
       return (f0', delDefs f0' {- this is dummy data, removing inner defs to avoid name conflicts -}, fj')
              
buildFunctions fn (ge@(DGVar g af)) v types defs =
    do gt <- getInputGraphType
       let typeVofG = getVertexType gt
           typeEofG = getEdgeType gt
           label = getName v
           label1 = getName g
           theType = lookupType label types
           typeG1 = lookupType label1 types
           -- the function body (copying the value from g)
           f0body = (expFieldAcc "prev" "v" [mkField label1,"_snd"])
           body = expBool True -- phase-termination = always true
       xx <- typingForBuildFunctions typeVofG typeEofG theType fn [f0body, body] [(label, theType)] [fieldStep fn]
       let [f0body', body'] = xx
           [fj', f0'] = buildDefVCtyped typeVofG typeEofG theType label [("_judge", typeDTBool, [], body'), ("_f0", theType, [], f0body')]   
       f0'' <- normalizeVertCompRecordOutputExp'' f0'
       return (f0'', delDefs f0'' {- this is dummy data, removing inner defs to avoid name conflicts -}, fj')

-- typing expressions in the given list bs
typingForBuildFunctions typeVofG typeEofG theType fn bs lts sls = 
    do let judgeVCEnv = buildJudgeVCEnv typeVofG typeEofG fn lts sls
           typing b = do b' <- typing2'' (judgeVCEnv) b theType
                         return b'
       mapM typing bs
  
typing2'' judgeVCEnv ftbody theType
    = do uid <- getUid
         let ((ftbody', _), uid') = typing2' (judgeVCEnv) ftbody theType uid
         setUid uid'
         return ftbody'
                
-- environment for typing judgment functions and so on.
-- lts is a list of labels used and their types, sls is a list of counters
buildJudgeVCEnv typeVofG typeEofG fn lts sls = 
    let testResType = newDataType fn
        extraTB = map (\(l, t) -> (mkField l, typeFunction [testResType, typePair [typeDTBool, t]])) lts ++ map (\sl -> (sl, typeFunction [testResType, typeDTInt])) sls
    in buildVertCompEnv typeVofG typeEofG (DTypeVar "tX") testResType (extraTB ++ buildProgEnv typeVofG typeEofG initEnv)

-- to build functions for giter (termination conditions are dealt with by buildDefsBody)
buildFunctionsIter buildDefsBody fn f0 ft g v types defs =
    do gt <- getInputGraphType
       let typeVofG = getVertexType gt
           typeEofG = getEdgeType gt
           label = getName v
           -- the output graph of ft
           gx = case lookupBy getName (getName ft) defs of
                  (Just (DGDefGF (DDefGraphFun _ _ _ (DGVar go _) _) _)) -> go
                  otherwise -> error ("why doesn't ft has this from?")
           labelX = getName gx
           s = sigma fn label (getName g) types typeVofG typeEofG
           f0' = let (Just (DGDefVI (DDefVertInit f defs' e af) a)) = lookupBy getName (getName f0) defs in DDefVertComp f defs' e af
           theType = lookupType label types
           theTypeX = lookupType labelX types
           -- building the body of ft (in the sense of pregel)
           ftbody = (expFieldAcc "prev" "v" [mkField labelX,"_snd"]) -- copy the output graph
           judgeVCEnv = buildJudgeVCEnv typeVofG typeEofG fn [(labelX, theTypeX), (label, theType)] [mkStep label]
           s' = sigma' fn label types typeVofG typeEofG
       ftbody' <- typing2'' (judgeVCEnv) ftbody theType
       (defs', body') <- buildDefsBody label judgeVCEnv s'
       let [fj', ft'] = buildDefVCtyped typeVofG typeEofG theType label [("_judge", typeDTBool, defs', body'), ("_ft", theType, [], ftbody')]
           -- this has to be ConsApp (to work with the current NtoPregel.hs )
       ft'' <- normalizeVertCompRecordOutputExp'' ft'
       f0'' <- normalizeVertCompRecordOutputExp'' f0'
       return (applySigma s f0' typeVofG typeEofG, ft'', fj')

normalizeVertCompRecordOutputExp'' ft
    = do uid <- getUid
         rs <- getRecordsSpecs
         let (ft', uid') = normalizeVertCompRecordOutputExp' rs ft uid
         setUid uid'
         return ft'
                                                  
-- to build functions for fregel (termination conditions are dealt with by buildDefsBody)
buildFunctionsPregel buildDefsBody fn f0 ft g v types defs = 
    do gt <- getInputGraphType
       let typeVofG = getVertexType gt
           typeEofG = getEdgeType gt
           label = getName v
           s = sigma fn label (getName g) types typeVofG typeEofG
           ft' = let (Just (DGDefVC vc a)) = lookupBy getName (getName ft) defs in vc
           f0' = let (Just (DGDefVI (DDefVertInit f defs' e af) a)) = lookupBy getName (getName f0) defs in DDefVertComp f defs' e af
           -- building the enviroment for the VertexCompute being made
           testResType = newDataType fn
           theType = lookupType label types
           judgeVCEnv = buildJudgeVCEnv typeVofG typeEofG fn [(label, theType)] [fieldStep fn]
           s' = sigma' fn label types typeVofG typeEofG
       (defs', body') <- buildDefsBody fn label judgeVCEnv s'
        -- TODO: dependency data?
       let [fj'] = buildDefVCtyped typeVofG typeEofG theType label [("_judge", typeDTBool, defs', body')]   
       ft'' <- normalizeVertCompRecordOutputExp'' ft'
       f0'' <- normalizeVertCompRecordOutputExp'' f0'
       return  (applySigma s f0'' typeVofG typeEofG, applySigma s ft'' typeVofG typeEofG, fj')

-- for giter: looking the individual counter
buildDefsBodyFi label judgeVCEnv = buildDefsBodyF (mkStep label) label judgeVCEnv
buildDefsBodyIi e label judgeVCEnv = buildDefsBodyI e (mkStep label) label judgeVCEnv
buildDefsBodyUi e label judgeVCEnv = buildDefsBodyU e (mkStep label) label judgeVCEnv

-- for fregel: looking the common step counter
buildDefsBodyFp fn label judgeVCEnv = buildDefsBodyF (fieldStep fn) label judgeVCEnv
buildDefsBodyIp e fn label judgeVCEnv = buildDefsBodyI e (fieldStep fn) label judgeVCEnv
buildDefsBodyUp e fn label judgeVCEnv = buildDefsBodyU e (fieldStep fn) label judgeVCEnv

-- to build judgment function for Fix
-- buildDefsBodyF stepLabel label judgeVCEnv idp vth s = 
buildDefsBodyF stepLabel label judgeVCEnv s = 
    do let aggV = "agg_"++label
           aggBody = expAggr "and" (expFieldAcc "prev" "u" [mkField label,"_snd"] ^== expFieldAcc "curr" "u" [mkField label,"_snd"]) "tg" []
           body = (((expFieldAcc "prev" "v" [stepLabel]) ^> (expInt 0))) ^&& (expVar aggV)
           -- body = expCheckTerm aggV
       -- typing of the aggregator used in the body of VC
       aggBody' <- typing2'' (judgeVCEnv) aggBody typeDTBool
       -- typing of the body of VC
       body' <- typing2'' ([(aggV, typeDTBool)]++judgeVCEnv) body typeDTBool
       --building the VC
       let defAgg = DDefVar (DVar aggV (DASTData typeDTBool [])) [] aggBody' (DASTData typeDTBool [])
       return ([defAgg], body')
             
-- to build judgment function for Iter
buildDefsBodyI e stepLabel label judgeVCEnv s =
    do let iterV = "iter_"++label
           body = (((expFieldAcc "prev" "v" [stepLabel]) ^== (expVar iterV))) --TODO: step -> step_label
       -- typing of the body of VC
       body' <- typing2'' ([(iterV, typeDTInt)]++judgeVCEnv) body typeDTBool
       --building the VC
       let defV = DDefVar (DVar iterV (DASTData typeDTInt [])) [] e (DASTData typeDTInt [])
       return ([defV], body')

-- to build judgment function for Until
buildDefsBodyU e stepLabel label judgeVCEnv s =
    do uid <- getUid
       let e' = s e
           (e'', defs, uid') = aggExtractionE e' uid
       setUid uid'
       return (defs, e'')

buildDefVCtyped typeVofG typeEofG theType label ntdbs = 
    let vertType = typeVertex [typeVofG, typeEofG]
        tblType = typeFunction[typeVertex [(DTypeVar "tY"), typeEofG], theType]
        vcType = typeFunction [vertType, tblType, tblType, typeDTBool]
        vcType' = typeFunction [vertType, tblType, tblType, theType]
        f (n, t, ds, b) = let vcType = typeFunction [vertType, tblType, tblType, t]
                          in defVCtyped (label++n) vcType ds b
    in map f ntdbs

lookupType label types = case lookup label types of (Just x) -> x; Nothing -> error ("why? " ++ label ++ " is not found in " ++ show types)
defVCtyped f t ds b = DDefVertComp (DFun f (DASTData t [])) ds b (DASTData t [])


instance GraphVarExtractable (DDefGraphFun DASTData) where
    extractGV (DDefGraphFun f v defs e a) = 
        do newVarChain          -- starts a new chain
           addTypeInfo (getName v) (getVertexType $ typeOf (getData v)) -- add the type info. of the argument
           mapM_ extractGV defs -- build children's chain
           chain <- popVarChain -- finish the chain
           saveChain f chain    -- save the chain
         
------------ topological sort of bindings in DDefVertComp with taking 'curr' into account --------
--  by analyzing which field of curr depends on what name
--  shallow analysis: no analysis on subexpressions (fields of records) 
--  So, the follwoing two programs do the same computation, but
--   let ab = Pair 1 (curr v .^ pval .^ _fst) in Dat ab         is NG (ab depends on ab); 
--   let a = 1; b = (curr v .^ pval .^ _fst) in Dat (Pair a b)  is OK ( b depends on a ).

sortWithCurr :: [(String, [String])] -> [String] -> DDefVertComp DASTData -> DDefVertComp DASTData
sortWithCurr fs phaseDataPrefix (x@(DDefVertComp f defs e a)) = res
  where (defs', ok) = topologicalSort (zip3 [0..] deps defs)
        names = map getNames defs
        deps  = map (genDeps (names, rdeps, phaseDataPrefix)) defs
        rdeps = genRecordDeps' fs names e
        res = if ok then DDefVertComp f (map snd defs') e a
              else error ("circular dependency through curr!\n" ++ unlines (ppDefVertComp x))

genRecordDeps :: [DRecordSpec a] -> [[String]] -> DExpr DASTData -> [([String], [Int])]
genRecordDeps rs ns e = genRecordDeps' (("Pair", ["_fst", "_snd"]):map extractFields rs) ns e

--union x y = nub (x++y)
unions = foldr union []
flatDeps depss = [([], unions (concatMap (map snd) depss))] 
addDeps deps = map (\(fs, ds) -> (fs, union ds deps))

mergeRDs dt de = union dt de -- correct?

genRecordDeps' :: [(String, [String])] -> [[String]] -> DExpr DASTData -> [([String], [Int])]
genRecordDeps' fs ns e = let ret = rec e in {- trace (show ret) $ -} ret
  where
    rec (DIf p t e _) = 
        let [([], deps)] = rec p
            dt = rec t
            de = rec e
            dte = mergeRDs dt de
        in addDeps deps dte -- adding the condition's deps
    rec (DTuple es _) = flatDeps $ map rec es -- conservative: flatterns the dependencies
    rec (DFunAp f es _) = addDeps (findIndex2 ns (getName f)) $ flatDeps $ map rec es -- conservative: flatterns the dependencies
    rec (DConsAp (DConstructor c _) es _) = 
        let Just flds = lookup c fs
            des = map rec es
        in concat $ zipWith (\field de -> map (\(fields, deps) -> (field:fields, deps)) de) flds des -- adding the fields of the record 
    rec (DFieldAcc (DCurr _ a) fs _) = error "Do not use curr in the last expression of VertexCmopute"
    rec (DFieldAcc _ _ _) = []
    rec (DFieldAccE _ _ _) = []
    rec (DAggr _ e _ es _) = flatDeps $ map rec (e:es) -- conservative: flatterns the dependencies
    rec (DVExp v _) = [([], findIndex2 ns (getName v))]
    rec (DCExp _ _) = []


class GenDepsable a where
  genDeps :: ([[String]], [([String], [Int])], [String]) -> a -> [Int]

getRdeps (_, rdeps, _) = rdeps
getENames (names, _, _) = names
getPDP (_, _, pdp) = pdp

instance (GenDepsable (DSmplDef a)) where
    genDeps env (DDefFun _ _ [] e _) = genDeps env e
    genDeps env (DDefVar _   [] e _) = genDeps env e
    genDeps env (DDefTuple _ [] e _) = genDeps env e

genDepss env = foldr (\a r -> nub (genDeps env a ++ r)) []

instance (GenDepsable (DExpr a)) where
    genDeps env (DIf p t e _) = nub (genDeps env p ++ (nub (genDeps env t ++ genDeps env e)))
    genDeps env (DTuple es _) = genDepss env es
    genDeps env (DFunAp f es _) = nub (genDeps env f ++ genDepss env es)
    genDeps env (DConsAp _ es _) = (genDepss env es)
    genDeps env (DFieldAcc (DCurr _ a) fs _) = 
        let f1:f2:sfs = map getName fs
            -- if [f1, f2] is of the fields that were added during the normalization process to indicate the phase, we have to add dependency. otherwise, no need to add dependency because the order of computation is guaranteed outside the VertexCompute being processed.
            rdeps = getRdeps env
            res = findIndex (\(sfs', _) -> isPrefixOf sfs' sfs) rdeps
        in if [f1, f2] == (getPDP env) 
           then case res of 
                  Just i -> snd (rdeps!!i)
                  Nothing -> error ("Something wrong! not found " ++ show sfs ++ " in " ++ show rdeps)
           else []
  -- add dependency of the field in the result
    genDeps env (DFieldAcc _ _ _) = []
    genDeps env (DFieldAccE _ _ _) = []
    genDeps env (DAggr _ e _ es _) = (genDepss env (e:es))
    genDeps env (DVExp v _) = genDeps env v
    genDeps env (DCExp _ _) = []

instance (GenDepsable (DVar a)) where
    genDeps env (DVar v _) = findIndex2 (getENames env) v

findIndex2 names v = maybe [] (\i -> [i]) (findIndex (\ns -> maybe False (const True) (findIndex (==v) ns)) names)

instance (GenDepsable (DFun a)) where
    genDeps env (DFun f _) = findIndex2 (getENames env) f
    genDeps env (DBinOp f _) = findIndex2 (getENames env) f

------------------------- driver ------------------------

-- this executes the big-step normalization as well as typechecking and the other transformations
normalizationAll ::Show a =>  DProgramSpec a -> DNormalized DASTData
normalizationAll p = makeNormalizedData (normalization p defaultOpt)

normalizationAll' ::Show a =>  DProgramSpec a -> Option -> DNormalized DASTData
normalizationAll' p ops = makeNormalizedData (normalization p ops)

-- compilation (typechecking and normalization) chain run on the AST
normalization :: Show a => DProgramSpec a -> Option -> (DProgramSpec DASTData, DUnique)
normalization ast ops = normalizationX numNormStages ast ops

-- TODO: K-normlization for graph expressions (at an early stage). currently, the parser does reject.

numNormStages = 11
normalizationX :: Show a => Int -> DProgramSpec a -> Option -> (DProgramSpec DASTData, DUnique)
normalizationX n ast0 ops = 
    let uid0 = 1024
        (ast1, uid1) = dependencyC uid0 ast0             -- 1. alpha renaming to unique names, topological sort by dependencies
        ast2 = typing ast1                               -- 2. typechecking
        ast3 = mapData (\(t,(d,_)) -> DASTData t d) ast2 -- 3. restructuring aux. data in the AST
        (ast4, uid4) = runInlining ast3 uid1             -- 4. inlining functions (dependency info. is broken here to make the process simple)
        (ast5, uid5) = aggExtraction ast4 uid4           -- 5. extraction of aggregation expressions
        ast6 = fixInputGraphType ast5                    -- 6. fix the type of input graph
        (ast7, uid7) = runTypeInstantiation ast6 uid5    -- 7. instantiating polymorphic funcs.
        (ast8, uid8) = makeOutputGraphVar ast7 uid7      -- 8. make the output expression to be a graph variable
        (ast9, uid9) = normalizeVertCompRecordOutputExp ast8 uid8 -- 9. make the output expression of a vertex computation function to be a simple constructor application
        astA = computeDependency ast9                    -- 10. updating the dependency info.
        astB = if zipOpt ops then moveGZipsToGMaps astA  -- 11. moving gzips into gmaps (and gzips)
               else astA  
    in ([ast3, ast4, ast5, ast6, ast7, ast8, ast9, astA, astB]!!(n-3), uid9)

-- moving gzips to gmaps so to remove redundant intermediate data about pairing
-- if g = gzip g1 g2 is used only inside a gmap, expand it.
-- assumption: the last expression is a graph variable
-- assumption: k-normalization about graph expressions
moveGZipsToGMaps :: DProgramSpec DASTData -> DProgramSpec DASTData
moveGZipsToGMaps (DProgramSpec rs (DProg f defs e app) ap) = 
          let (dgfs, dgvs, ds) = splitGDefs defs
              ginType = getInputType $ getType f
              g = DVar "g" (DASTData (ginType) [])              -- the input is 'g'
              dgvs' = map (\dgv -> DGDefGV dgv (getData dgv)) $ moveZM g (map (\(DGDefGV dgv a) -> dgv) dgvs)  -- moving in the outermost level
              dgfs' = map moveZMinGF dgfs           -- moving in other graph functions
              defs' = ds ++ dgfs' ++ dgvs'
          in DProgramSpec rs (DProg f defs' e app) ap

-- wrapper for moveZM
moveZMinGF (DGDefGF (DDefGraphFun f gin dgvs e a) ap)
           = let dgvs' = moveZM gin dgvs
             in DGDefGF (DDefGraphFun f gin dgvs' e a) ap

-- moving single-use gzips inside gmaps
-- gin is the input graph
-- gvs is a list of graph variable definitions whose rhss are k-normal (no nested expressions).
moveZM :: DVar DASTData -> [DDefGraphVar DASTData] -> [DDefGraphVar DASTData]
moveZM gin gvs = run gvs
    where
      run [] = []
      run (gv:gvs)
          = case gv of 
              (DDefGraphVar gn (gez@(DGZip ge1 ge2 a)) a')
                  -> let gnn = getName gn
                         gused = filter (isUsed gnn) gvs
                     in if length gused == 1 && isMapOrZip (head gused)
                        then run (map (expandZip gnn gez) gvs)
                        else gv:run gvs
              otherwise -> gv:run gvs
-- replace variable gnn with gez in ge
expandZip :: String -> DGraphExpr DASTData -> DDefGraphVar DASTData -> DDefGraphVar DASTData
expandZip gnn gez (DDefGraphVar gv ge a) = (DDefGraphVar gv (run ge) a)
    where
      run (DPregel fi fs tc ge a) = (DPregel fi fs tc (run ge) a)
      run (DGMap f ge a) = (DGMap f (run ge) a)
      run (DGZip ge1 ge2 a) = (DGZip (run ge1) (run ge2) a)
      run (DGIter fi fs tc ge a) = (DGIter fi fs tc (run ge) a)
      run (ge@(DGVar (DVar n a) ap)) = if gnn == n then gez else ge
  
                           
isMapOrZip (DDefGraphVar gv (DGMap f ge a) a') = True    
isMapOrZip (DDefGraphVar gv (DGZip ge1 ge2 a) a') = True
isMapOrZip _ = False
                                   
isUsed gn (DDefGraphVar gv ge a) = isUsed' ge
    where 
      isUsed' (DPregel fi fs tc ge a) = isUsed' ge
      isUsed' (DGMap f ge a) = isUsed' ge
      isUsed' (DGZip ge1 ge2 a) = isUsed' ge1 || isUsed' ge2
      isUsed' (DGIter fi fs tc ge a) = isUsed' ge
      isUsed' (DGVar (DVar n a) ap) = gn == n
                            
-- dividing a list of definitions into three kinds: graph functions, graph variables, and others
splitGDefs defs = splitGDefs' [] [] [] defs
splitGDefs' dgfs dgvs ds [] = (reverse dgfs, reverse dgvs, reverse ds)
splitGDefs' dgfs dgvs ds ((d@(DGDefGF _ _)):defs) = splitGDefs' (d:dgfs) dgvs ds defs
splitGDefs' dgfs dgvs ds ((d@(DGDefGV _ _)):defs) = splitGDefs' dgfs (d:dgvs) ds defs
splitGDefs' dgfs dgvs ds (d:defs) = splitGDefs' dgfs dgvs (d:ds) defs

       
-- make the last expression to be a variable binding a graph (8th step)
makeOutputGraphVar :: DProgramSpec DASTData -> DUnique -> (DProgramSpec DASTData, DUnique)
makeOutputGraphVar p uid = 
    case p of 
      (DProgramSpec rs (DProg f defs (DGVar _ _) app) ap) -> (p, uid)
      (DProgramSpec rs (DProg f defs e app) ap) -> 
          let ae = (getData e)
              (nv, uid') = genNewName uid "g_v"
              gv = (DVar nv ae)
              def = DGDefGV (DDefGraphVar gv e ae) ae
              defs' = defs++[def]
              e' = DGVar gv ae
          in (DProgramSpec rs (DProg f defs' e' app) ap, uid')

-- fix the input graph type (stupid) (6th step)
fixInputGraphType' :: DTypeInfo -> DTypeInfo -> DProgramSpec DASTData -> DProgramSpec DASTData
fixInputGraphType' tv te p = let gt = typeGraph [tv, te]
                                 pt = getType p
                                 pgt = getInputType pt
                                 s = unify [(gt, pgt)]
                                 p' = mapData (\a -> a {typeOf = apply s (typeOf a)}) p
                             in p'

fixInputGraphType :: DProgramSpec DASTData -> DProgramSpec DASTData
fixInputGraphType p = let pgt = getInputType (getType p)
                      in case pgt of 
                              DTypeTerm "Graph" [DTypeVar _, DTypeVar _] -> fixInputGraphType' typeNull typeNull p
                              DTypeTerm "Graph" [DTypeVar _, et]         -> fixInputGraphType' typeNull et p
                              DTypeTerm "Graph" [vt        , DTypeVar _] -> fixInputGraphType' vt typeNull p
                              DTypeTerm "Graph" [vt        , et        ] ->  p


------------------------- pretty printer with types ------------------

ppprint = putStrLn . unlines . ppProgramSpec -- without types (;

ppprintStr = unlines . ppProgramSpec
ppprintStr' n p = (getPName p ++ "__N" ++ show n ++ ".fgl", ppprintStr p)

pptprintStr :: PrettyPrinterWithType a => a -> String
pptprintStr = unlines . ppt
pptprintStr' n p = (getPName p ++ "__N" ++ show n ++ ".tfgl", pptprintStr p)

                   
getPName (DProgramSpec rs (DProg f defs _ _) a) = getName f        

pptprint :: PrettyPrinterWithType a => a -> IO ()
pptprint = putStrLn . unlines . ppt

type Misc = DASTData

class PrettyPrinterWithType a where
    ppt :: a -> [String]

instance PrettyPrinterWithType (DProgramSpec Misc) where
    ppt (DProgramSpec rs p _) = concatMap ppt rs ++ ppt p

instance PrettyPrinterWithType (DRecordSpec Misc) where
    ppt r = ppRecordSpec r

instance PrettyPrinterWithType (DProg Misc) where
    ppt (x@(DProg f ds e _)) = [typeSig f ++ dep x, ppFun f ++ " g = "] ++ rest
        where rest = if length ds == 0 then indent (ppt e)
                     else indent $ let_in (concat $ insList ";" "" $ map ppt ds) (ppt e)

typeSig f = ppFun f ++ "::" ++ prettyShow (getType f)
typeSigV v = ppVar v ++ "::" ++ prettyShow (getType v)
typeSigC c = ppConst c ++ "::" ++ prettyShow (getType c)
typeSigCr c = ppConstructor c ++ "::" ++ prettyShow (getType c)

getDepNames x = depOf $ getData x

dep x = "  -- depends on " ++ ppList "," (getDepNames x)

instance PrettyPrinterWithType (DGroundDef Misc) where
    ppt (DGDefVC d _) = ppt d
    ppt (DGDefVI d _) = ppt d
    ppt (DGDefGV d _) = ppt d 
    ppt (DGDefGF d _) = ppt d
    ppt (DGDefSmpl d _) = ppt d

instance PrettyPrinterWithType (DDefVertComp Misc) where
    ppt (x@(DDefVertComp f ds e _)) =
        if length ds == 0 then [typeSig f ++ dep x] ++ indentWith (ppFun f ++ " v prev curr = ") (ppt e)
        else [typeSig f ++ dep x] ++ [ppFun f ++" v prev curr = "] ++ indent (let_in (concat $ insList ";" "" $map ppt ds) (ppt e))

instance PrettyPrinterWithType (DDefVertInit Misc) where
    ppt (x@(DDefVertInit f ds e _)) =
        if length ds == 0 then [typeSig f ++ dep x] ++ indentWith (ppFun f ++ " v = ") (ppt e)
        else  [typeSig f ++ dep x] ++ [ppFun f ++" v = "] ++ indent (let_in (concat $ insList ";" "" $map ppt ds) (ppt e))

instance PrettyPrinterWithType (DDefGraphVar Misc) where
    ppt (x@(DDefGraphVar v e _)) =
        [typeSigV v ++ dep x] ++ indentWith (ppVar v ++ " = ") (ppt e)
  

instance PrettyPrinterWithType (DDefGraphFun Misc) where
    ppt (x@(DDefGraphFun f v ds e _)) =
        if length ds == 0
        then header ++ indent (ppt e)
        else header ++ indent (let_in (concat $ insList ";" "" $ map ppt ds) (ppt e))
            where header = [ typeSig f ++ dep x, ppFun f ++ " " ++ (ppVar v) ++ " = " ]

instance PrettyPrinterWithType (DGraphExpr Misc) where
    ppt (x@(DPregel f0 ft t g _)) =
        ["("++(ppList " " $ ["pregel", ppFun f0, ppFun ft] ++ ppt t ++ ppt g) ++ "::" ++ prettyShow (getType x)++")"]
    ppt (x@(DGMap f g _)) =
        ["("++ppList " "  (["gmap", ppFun f] ++ ppt g) ++ "::" ++ prettyShow (getType x)++")"]
    ppt (x@(DGZip g1 g2 _)) =
        ["("++ppList " "  (["gzip"] ++ ppt g1 ++ ppt g2) ++ "::" ++ prettyShow (getType x)++")"]
    ppt (x@(DGIter f0 ft t g _)) =
        ["("++(ppList " " $ ["iter", ppFun f0, ppFun ft] ++ ppt t ++ ppt g) ++ "::" ++ prettyShow (getType x)++")"]
    ppt (x@(DGVar g _)) = 
        ["("++ppVar g ++ "::" ++ prettyShow (getType x)++")"]

instance PrettyPrinterWithType (DTermination Misc) where
    ppt (DTermF _)= ["Fixpoint"]
    ppt (DTermI e _) = ["(Iter (" ++ ppList " " (ppt e) ++ "))"]
    ppt (DTermU e _) = ["(Until (\\g->" ++ ppList " " (ppt e) ++ "))"]

instance PrettyPrinterWithType (DVar Misc) where
    ppt x = [ppVar x ++ "::" ++ prettyShow (getType x) ++ dep x]

instance PrettyPrinterWithType (DSmplDef Misc) where
    ppt (x@(DDefFun f vs ds e _)) =
        if length ds == 0
        then header ++ indent (ppt e)
        else header ++ indent (let_in (concatMap ppt ds) (ppt e))
            where header = [typeSig f ++ dep x, ppFun f ++ " " ++ ppList " " (map ppVar vs) ++ " = " ]
    ppt (x@(DDefVar v ds e _)) = 
        if length ds == 0
        then header ++ indent (ppt e)
        else header ++ indent (let_in (concatMap ppt ds) (ppt e))
            where header = [typeSigV v ++ dep x, ppVar v ++ " = "]
    ppt (x@(DDefTuple vs ds e _)) = 
        if length ds == 0
        then header ++ indent (ppt e)
        else header ++ indent (let_in (concatMap ppt ds) (ppt e))
            where header = [tt ++ "::" ++ prettyShow (getType x) ++ dep x, tt ++" = "]
                  tt = "(" ++ ppList ", " (map ppVar vs) ++ ")"

enclose' xs = indentWith "(" xs'
    where xs' = init xs ++ [last xs++")"]

addSig x xs = enclose' xs''
 where xs' = enclose' xs
       xs'' = init xs' ++ [last xs'++"::"++ prettyShow (getType x)]

instance PrettyPrinterWithType (DExpr Misc) where
    ppt (x@(DIf c t e _)) = addSig x (indentWith "if " (ppt c) ++
                                      indentWith "then " (ppt t) ++
                                      indentWith "else " (ppt e))
    ppt (x@(DVExp v _)) = ["("++typeSigV v++")"]
    ppt (x@(DCExp c _)) = ["("++typeSigC c++")"]
    ppt (x@(DFunAp f es _)) = addSig x [ppList " " (ppFun f: map flatE' es)]
    ppt (x@(DConsAp c es _)) = addSig x [ppList " " (ppConstructor c: map flatE' es)]
    ppt (x@(DTuple es _)) = addSig x $ indentWith "(" (concat $ insList "," ")" $ map ppt es)
    ppt (x@(DFieldAcc e fs _)) = addSig x $ [ ppTableExpr e ++ (concat $ map ((".^"++).ppField) fs) ]
    ppt (x@(DFieldAccE e fs _)) = addSig x $ [ ppEdge e ++ (concat $ map ((".^"++).ppField) fs) ]
    ppt (x@(DAggr a e g es _)) =
        addSig x $ [ppAgg a ++ " [ " ++ flatE2' e ++  " | " ++ ppGen g ++ ps ++ " ] "]
            where ps = if length es == 0 then ""
                       else "," ++ ppList "," (map flatE' es)
    --ppt e = ppExpr e

flatE' :: (PrettyPrinterWithType (DExpr a)) => DExpr a -> String
flatE' = {- enclose . -} ppList " " . ppt

flatE2' :: (PrettyPrinterWithType (DExpr a)) => DExpr a -> String
flatE2' = ppList " " . ppt


printN' (DNormalized fn is rs nd ps deps [ip] defs) =
   putStrLn(unlines $ [
            "fname: " ++ fn,
            "is: " ++ show is,
            "rs: "] ++ concatMap (indentWith "  " . ppRecordSpec) rs ++
            indentWith "  "  (ppRecordSpec (fst nd)) ++
              [
            "ip: " ++ show ip,
            "ps: "] ++
            concatMap (\(is, str, fi, fs, fj) -> [" id: " ++ show is] ++ [ " nm: " ++ str] ++ indentWith " fi:" (ppDefVertComp fi) ++ indentWith " fs:" (ppDefVertComp fs) ++ indentWith " fj:" (ppDefVertComp fj)) ps
            ++ ["defs: "] ++ concatMap ppSmplDef defs
            )

------------ pretty printer for normalized oen (n_fregel) ----------------
printNStr = unlines . ppN

getFName (DNormalized fn is rs nd ps deps [ip] defs) = fn

printNStr' p = (getFName p ++ "__N" ++ show (numNormStages - 1) ++  ".fgl", printNStr p) 

-- using Int indtead of Maybe Phase for simplicity
printN = putStrLn . printNStr

addSemicolon :: [[String]] -> [[String]]
addSemicolon (x:[]) = x:[]
addSemicolon (x:xs) = addSemicolonLast x:addSemicolon xs
                      where addSemicolonLast (x:[]) = (x++";"):[]
                            addSemicolonLast (x:xs) = x:addSemicolonLast xs
                                                      
ppN (DNormalized fn is rs' nd ps deps [ip] defs) =
  let rs = fst nd:rs'
      (Just theRecord) = lookupBy (\(DRecordSpec c _ _) -> getName c) (newDataTypeName fn) rs
  in concatMap ppRecordSpec rs ++ [fn ++ " g = "] ++ indent (indentWith "let " (concat $ addSemicolon (map ppSmplDef defs ++ map (genFtX fn is) ps ++ [genFt fn ps is deps theRecord] ++ [genF0 fn ip rs theRecord])) ++ ["in fregel f0 ft Fix g"])

genF0 fn ip rs (DRecordSpec _ (_:_:fts) _) = ["f0 v = " ++ newDataTypeName fn ++" " ++ show ip ++ " 0" ++ concatMap (" "++) (map (genInitData.snd) fts)]
  where
    genInitData (DTInt _) = "0"
    genInitData (DTBool _) = "False"
    genInitData (DTString _) = "\"\""
    genInitData (DTDouble _) = "0.0"
    genInitData (DTTuple ts _) = "(" ++ ppList ", " (map genInitData ts) ++ ")"
    genInitData (DTRecord (DConstructor "Pair" _) [t1,t2] _) = "(Pair " ++ genInitData t1 ++ " " ++ genInitData t2 ++ ")"
    genInitData (DTRecord c ts _) =
      let (Just (DRecordSpec _ fts _)) = lookupBy (\(DRecordSpec c _ _) -> getName c) (getName c) rs
      in "(" ++ (getName c) ++ concatMap (" "++) (map (genInitData.snd) fts) ++ ")"

genFt fn ps is deps theRecord = ["ft v prev curr = "] ++ indentWith "  let " (concat $ addSemicolon (map (genFt0 fn) ps ++ [genEnd fn ps] ++ [genStep fn ps] ++ [genPhase fn ps deps] ++ map (genStepX fn ps) is))  ++ ["  in " ++ genFtBody fn ps is theRecord]
genFt0 :: DVarName -> (DPhaseID, String, DDefVertComp DASTData, DDefVertComp DASTData, DDefVertComp DASTData) -> [String]
genFt0 fn (pid, label, f0, ft, fj) = 
  ppSmplDef $ defVar ("d_"++label) [] (expIf (expFieldAcc "prev" "v" [(fieldPhase fn)] ^== (expInt pid)) (expFun ("ft_"++label) [expVar "v", expVar "prev", expVar "curr"]) (expFieldAcc "prev" "v" [mkField label]))

genStepX fn ps pid = ppSmplDef $ defVar ("s_" ++ label) [] (expIf (expFieldAcc "prev" "v" [(fieldPhase fn)] ^== (expInt pid)) (expIf (expFun "not" [expVar "p_end"]) (expBin "+" (expInt 1) (expFieldAcc "prev" "v" [mkStep label])) (expInt 0)) (expFieldAcc "prev" "v" [mkStep label]))
  where label = lookupLabel pid ps


genPhase fn ps deps = concat $ addSemicolon $ map ppSmplDef [defNext, defStay, defPhase]
  where
    defStay = defFun "stay" ["pid"] [] (recs 0 deps)
    defNext = defFun "next" ["pid"] [] (recs 1 deps)
    recs k [] = expInt (maxpid ps)
    recs k ((pid, fps):ds) = expIf (expVar "pid" ^== expInt pid) (expInt ((if length fps > k then snd (fps!!k) else maxpid ps))) (recs k ds)
    defPhase = defVar "phase'" [] (expIf (expVar "p_end") (expFun ("next") [(expFieldAcc "prev" "v" [(fieldPhase fn)])]) (expFun ("stay") [(expFieldAcc "prev" "v" [(fieldPhase fn)])]))

genFtBody fn ps is (DRecordSpec _ (_:_:fts) _) = newDataTypeName fn ++ " phase' step'" ++ concatMap (" "++) (map (\(_, label, _,_,_) -> "d_" ++ label) ps) ++ concatMap (" "++) (map (\i -> "s_" ++ lookupLabel i ps) is)

lookupLabel i ps = let (Just l) = lookup i (map (\(i,l,_,_,_) -> (i, l)) ps) in l

genEnd fn ps = ppSmplDef $ defVar "p_end" [] (genEndrec ps)
  where
    genEndrec [] = expBool False
    genEndrec ((pid, label, f0, ft, fj):ps) = expIf (expFieldAcc "prev" "v" [(fieldPhase fn)] ^== (expInt pid)) (expFieldAcc "curr" "v" [mkField label, "_fst"]) (genEndrec ps)

genStep fn ps = ppSmplDef $ defVar "step'" [] (expIf ((expFun "not" [expFieldAcc "prev" "v" [(fieldPhase fn)] ^== (expInt (maxpid ps))]) ^&& (expFieldAcc "prev" "v" [(fieldPhase fn)] ^== expFieldAcc "curr" "v" [(fieldPhase fn)] )) (expBin "+" (expInt 1) (expFieldAcc "prev" "v" [(fieldStep fn)])) (expInt 0))
maxpid ps = 1 + maximum (map (\(pid,_,_,_,_) -> pid) ps)
       
getDefsOfVC (DDefVertComp _ defs _ _) = defs

-- omitting type info.
genFtX :: DVarName -> [DPhaseID] -> (DPhaseID, String, DDefVertComp DASTData, DDefVertComp DASTData, DDefVertComp DASTData) -> [String]
genFtX fn is (pid, label, f0, ft, fj) = 
  let (DDefVertComp _ defs0 e0 _) = mapData (\_ -> "") f0
      (DDefVertComp _ defst et _) = mapData (\_ -> "") ft
      (DDefVertComp _ defsj ej _) = mapData (\_ -> "") fj
      -- if this is iter, then look at its step counter
      st = if (elem pid is) then mkStep label else (fieldStep fn) 
      defD = defVar "d" [] (expIf (expFieldAcc "prev" "v" [st] ^== (expInt 0)) e0 et)
      defEnd = defVar "end" [] ej
  in ppDefVertComp $ defVertComp ("ft_"++label) (defs0++defst++[defD]++defsj++[defEnd]) (expConstructor "Pair" [(expVar "end"),(expVar "d")])



     
---------------- for flatterning nested records (not used yet) ------------------------
genFieldListsOfType :: [DRecordSpec a] -> DTypeInfo -> [([String], DTypeInfo)]
genFieldListsOfType rs t = genFieldListsOfType' (buildExpander rs) t


genFieldListsOfType' :: (DTypeInfo -> [(String, DTypeInfo)]) -> DTypeInfo -> [([String], DTypeInfo)]
genFieldListsOfType' expander t = case (expander t) of 
                                  [] -> [([], t)]
                                  fts -> concatMap f fts 
  where f (field, t) = map (\(fields, t') -> (field:fields, t')) (genFieldListsOfType' expander t)


buildExpander :: [DRecordSpec a] -> DTypeInfo -> [(String, DTypeInfo)]
buildExpander rs (DTypeTerm tn ts) = 
    if tn == "Pair" then [("_fst", head ts), ("_snd", head (tail ts))]
    else if isPrimitive tn then [] 
         else case findRecord rs tn of
                Just fts -> fts
                Nothing -> error ("something wrong. not found " ++ tn)

findRecord [] _ = Nothing
findRecord (DRecordSpec (DConstructor c _) fts _:rs) tn =
    if c == tn then Just (map (\(DField f _, t) -> (f, toTypeInfo t)) fts) 
    else findRecord rs tn

toTypeInfo (DTInt _) = typeDTInt
toTypeInfo (DTBool _) = typeDTBool
toTypeInfo (DTString _) = typeDTString
toTypeInfo (DTDouble _) = typeDTDouble