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

{-

Given a vertex-compute function, this makes its output expression to be of the form RecordConstructor e1 e2 ... ek (9th step)

Assumption: 

  - typed.
  - inlined.

-}

module VCOutputNormalization where

import Spec


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

import ASTData
import TypeChecker(typing, unify, apply, buildVertCompEnv, buildProgEnv, typeCheck, typingTypeExpression)

type DONState = DUnique


normalizeVertCompRecordOutputExp' :: [DRecordSpec DASTData] -> DDefVertComp DASTData -> DUnique -> (DDefVertComp DASTData, DUnique)
normalizeVertCompRecordOutputExp' rs x uid
 = runState act uid
   where act = do x' <- normVCOutput' (rs++[nullRecordSpec]) x
                  return x'

normVCOutput' rs (DDefVertComp f defs e a)
 = do let rt = (getTypeName (getType e))
          r = maybe (error ("unknown record type: " ++ rt)) id $ lookupBy (getName) rt rs
      (defs', e') <- normVCOutputExp r e
      return (DDefVertComp f (defs++defs') e' a)

normalizeVertCompRecordOutputExp :: DProgramSpec DASTData -> DUnique -> (DProgramSpec DASTData, DUnique)
normalizeVertCompRecordOutputExp (DProgramSpec rs (DProg f defs e a) a') uid
 = (DProgramSpec rs (DProg f defs' e a) a', uid')
   where (defs', uid') = runState act uid 
         act = do defs' <- mapM (normVCOutput (rs++[nullRecordSpec])) defs
                  return defs'
normVCOutput :: [DRecordSpec DASTData] -> DGroundDef DASTData -> State DONState  (DGroundDef DASTData)
normVCOutput rs (DGDefVC (DDefVertComp f defs e a) a')
 = do let rt = (getTypeName (getType e))
          r = maybe (error ("unknown record type: " ++ rt)) id $ lookupBy (getName) rt rs
      (defs', e') <- normVCOutputExp r e
      return (DGDefVC (DDefVertComp f (defs++defs') e' a) a')
normVCOutput rs (DGDefVI (DDefVertInit f defs e a) a')
 = do let rt = (getTypeName (getType e))
          r = maybe (error ("unknown record type: " ++ rt)) id $ lookupBy (getName) rt rs
      (defs', e') <- normVCOutputExp r e
      return (DGDefVI (DDefVertInit f (defs++defs') e' a) a')
normVCOutput rs x = do return x


--make the final expression to be a construcgtor application
normVCOutputExp :: DRecordSpec DASTData -> DExpr DASTData -> State DONState ([DSmplDef DASTData], DExpr DASTData)
normVCOutputExp r (e@(DConsAp c es a)) = do splitToDefs e
normVCOutputExp (r@(DRecordSpec c fts a)) e = 
    do let es = splitByFields r e
       splitToDefs (DConsAp c es (getData e))

getNewName :: String -> State (DONState) DVarName
getNewName s = 
    do (i) <- get 
       let (n, i') = genNewName i s
       put (i')
       return n

splitToDefs :: DExpr DASTData -> State DONState ([DSmplDef DASTData], DExpr DASTData)
splitToDefs (DConsAp c es a) = 
         -- if a code generator requries that no operand of the constructor application can be a constant expression, then remove this branch.
    do vs <- mapM (\e -> case e of DVExp _ _ -> return Nothing; DCExp _ _ -> return Nothing; otherwise -> (do n <- (getNewName "var"); return (Just n))) es
       let ves = (zip vs es)
           defss = map (\(v,e) -> maybe [] (\n -> [DDefVar (DVar n (getData e)) [] e (getData e)]) v) ves
       return  (concat defss, DConsAp c (map (\(x, e) -> case x of [] -> e; [DDefVar v _ _ a] -> DVExp v a) (zip defss es)) a)
       

splitByFields :: DRecordSpec DASTData -> DExpr DASTData -> [DExpr DASTData]
splitByFields (DRecordSpec c fts a) e = map (\i -> extractFieldExp e (fts!!i) i) [0..(length fts - 1)]

extractFieldExp :: DExpr DASTData -> (DField DASTData, DType DASTData) -> Int -> DExpr DASTData
extractFieldExp e (f, t) i = rec e -- this must be the i-th component of the e
  where t' = typingTypeExpression t
        rec (DIf e1 e2 e3 a) = let e2' = (rec e2) 
                                   e3' = (rec e3)
                               in DIf e1 e2' e3' (a {typeOf = getType e2'} )
        rec (DTuple es a) = error "something wrong!"
        rec (DFunAp f es a) = error "something wrong!"
        rec (DConsAp c es a) = es !! i    -- this must be correct
        rec (DFieldAcc te fs a) = DFieldAcc te (fs++[setType t' f]) (a {typeOf = t'})
        rec (DFieldAccE ee fs a) = DFieldAccE ee (fs++[setType t' f]) (a {typeOf = t'})
        rec (DAggr agg e g ps a) = error "something wrong!"
        rec (DVExp v a) = error "not supported yet"
        rec (DCExp c a) = error "something wrong!"